Repository: scikit-hep/uproot5

Issue:
Add support for current RNTuple files
Due to some recent changes in RNTuple's format, uproot is not able to read RNTuples correctly anymore.

A reproducer:

```
python -c 'import uproot; print(uproot.__version__); uproot.open("https://xrootd-local.unl.edu:1094//store/user/AGC/nanoaod-rntuple/zstd/TT_TuneCUETP8M1_13TeV-powheg-pythia8/cmsopendata2015_ttbar_19980_PU25nsData2015v1_76X_mcRun2_asymptotic_v12_ext3-v1_00000_0000.root")["Events"].arrays(["nTau"])
```
results in

```
5.0.10
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/blue/Repos/analysis-grand-challenge-iris-hep/analyses/cms-open-data-ttbar/venv-uproot5/lib/python3.11/site-packages/uproot/models/RNTuple.py", line 394, in arrays
    entry_stop = entry_stop or self._length
                               ^^^^^^^^^^^^
  File "/home/blue/Repos/analysis-grand-challenge-iris-hep/analyses/cms-open-data-ttbar/venv-uproot5/lib/python3.11/site-packages/uproot/models/RNTuple.py", line 180, in _length
    return sum(x.num_entries for x in self.cluster_summaries)
                                      ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/blue/Repos/analysis-grand-challenge-iris-hep/analyses/cms-open-data-ttbar/venv-uproot5/lib/python3.11/site-packages/uproot/models/RNTuple.py", line 175, in cluster_summaries
    return self.footer.cluster_summaries
           ^^^^^^^^^^^
  File "/home/blue/Repos/analysis-grand-challenge-iris-hep/analyses/cms-open-data-ttbar/venv-uproot5/lib/python3.11/site-packages/uproot/models/RNTuple.py", line 164, in footer
    f = FooterReader().read(self._footer_chunk, cursor, context)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/blue/Repos/analysis-grand-challenge-iris-hep/analyses/cms-open-data-ttbar/venv-uproot5/lib/python3.11/site-packages/uproot/models/RNTuple.py", line 697, in read
    out.extension_links = self.extension_header_links.read(chunk, cursor, context)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/blue/Repos/analysis-grand-challenge-iris-hep/analyses/cms-open-data-ttbar/venv-uproot5/lib/python3.11/site-packages/uproot/models/RNTuple.py", line 556, in read
    assert num_bytes < 0, f"num_bytes={num_bytes}"
           ^^^^^^^^^^^^^
AssertionError: num_bytes=36
```

IIUC the Julia implementation already implemented the necessary changes at https://github.com/JuliaHEP/UnROOT.jl/pull/264/files 

EDIT:
maybe I should have marked it as a bug?


Discussion:
@ioanaif has started to look into this, and found that some of the changes are headers/metadata, while another is that [variable-length integers](https://lucene.apache.org/core/3_5_0/fileformats.html#VInt) and [zig-zag encoding](https://protobuf.dev/programming-guides/encoding/#signed-ints) are now included in RNTuple.

Conversion to and from a variable-length integer format is not NumPy-vectorizable: it is necessary to write for loops to do this conversion. Here's what that looks like in Python:

```python
def to_varint(data):
    assert issubclass(data.dtype.type, np.uint64)

    output = []
    for value in data:
        mask = np.uint64(0x7f)
        more = np.uint64(np.iinfo(np.uint64).max)
        for shift in np.arange(0, 7 * 9, 7, dtype=np.uint64):
            byte = ((value & mask) >> shift).astype(np.uint8)
            mask <<= np.uint64(7)
            more <<= np.uint64(7)

            if not (value & more):
                output.append(byte)
                break
            else:
                output.append(byte | np.uint8(0x80))

    return b"".join(output)

def from_varint(buffer):
    data = []
    pos = 0
    while pos < len(buffer):
        shift = np.uint64(0)
        result = np.uint64(0)
        while True:
            byte = np.uint64(buffer[pos])
            pos += 1

            if shift == 7 * 9:
                raise Exception("number is too big for uint64")

            result |= (byte & np.uint64(0x7f)) << shift
            shift += np.uint64(7)

            if not (byte & np.uint64(0x80)):
                break

        data.append(result)

    return np.array(data)
```

```python
>>> data = np.array([0, 1, 2, 127, 128, 129, 130, 16383, 16384, 16385], np.uint64)
>>> data.tolist()
[0, 1, 2, 127, 128, 129, 130, 16383, 16384, 16385]

>>> buffer = to_varint(data)
>>> buffer
b'\x00\x01\x02\x7f\x80\x01\x81\x01\x82\x01\xff\x7f\x80\x80\x01\x81\x80\x01'

>>> from_varint(buffer).tolist()
[0, 1, 2, 127, 128, 129, 130, 16383, 16384, 16385]
```

But we want to avoid Python for loops. AwkwardForth provides a way to do that when converting _from_ a variable-length encoding into integers. Here's how it can be done:

```python
>>> from awkward.forth import ForthMachine64
>>> vm = ForthMachine64("""
... input buffer
... output data uint64
... 
... begin
...     buffer varint-> data
... again
... """)

>>> buffer = to_varint(np.array([0, 1, 2, 127, 128, 129, 130, 16383, 16384, 16385], np.uint64))

>>> vm.run({"buffer": buffer}, raise_read_beyond=False)
'read beyond'
>>> vm.outputs["data"]
array([    0,     1,     2,   127,   128,   129,   130, 16383, 16384,
       16385], dtype=uint64)
```

What the above does:

* declares [input and output buffers](https://awkward-array.org/doc/main/reference/awkwardforth.html#variables-inputs-and-outputs) (inputs are given; outputs are created and grow as needed)
* uses the [`varint->`](https://awkward-array.org/doc/main/reference/awkwardforth.html#variable-length-integers) word to decode one variable-length integer from the input `buffer` to the output `data`
* uses the [`begin .. again` construct](https://awkward-array.org/doc/main/reference/awkwardforth.html#begin-again) to do an infinite loop (like `while True`)
* uses `raise_read_beyond=False` to catch a "read beyond length of input buffer" exception and return it as a string instead of raising a Python exception.

But since AwkwardForth is compiled code, it's a lot faster than pure Python:

```python
%%timeit
from_varint(buffer);
# 5.11 s ± 15.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
vm.run({"buffer": buffer}, raise_read_beyond=False);
# 21.3 ms ± 114 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

(240× faster, in this case).

The value of this encoding is that integers that are close to zero use the fewest bytes, but any integer can be encoded (including integers larger than `uint64`, though that's all we care about when NumPy is involved).

You could do the same thing with signed integers, but small signed integers like `-1` would be encoded with the most bytes, because `-1` is `0xffffffffffffffff` in `int64`.

```python
>>> hex(np.int64(-1).view(np.uint64))
'0xffffffffffffffff'
```

So, on top of the variable-length encoding, signed integers are zig-zag encoded, mapping e.g. `[0, -1, 1, -2, ...]` as `[0, 1, 2, 3, ...]`. That step can be vectorized:

https://github.com/JuliaHEP/UnROOT.jl/blob/d50081090d95b44138098e25b4102e0d01f270a6/src/RNTuple/fieldcolumn_reading.jl#L87-L88

or

```python
from_zigzag = lambda n: (n >> 1) ^ -(n & 1)
to_zigzag = lambda n: (n << 1) ^ (n >> 63)
```

in Python. However, AwkwardForth has a built-in word for it, [`zigzag->`](https://awkward-array.org/doc/main/reference/awkwardforth.html#variable-length-integers) (which does both the zig-zag and the variable-length decoding), so you can just use that.

```python
>>> vm = ForthMachine64("""
... input buffer
... output data int64
... 
... begin
...     buffer zigzag-> data
... again
... """)

>>> buffer = to_varint(
...     to_zigzag(np.array([0, -1, 1, -2, 2, 100, -100, 1000, -1000], np.int64)).astype(np.uint64)
... )

>>> vm.run({"buffer": buffer}, raise_read_beyond=False)
'read beyond'
>>> vm.outputs["data"]
array([    0,    -1,     1,    -2,     2,   100,  -100,  1000, -1000])
```
> @ioanaif has started to look into this, and found that some of the changes are headers/metadata, while another is that [variable-length integers](https://lucene.apache.org/core/3_5_0/fileformats.html#VInt) and [zig-zag encoding](https://protobuf.dev/programming-guides/encoding/#signed-ints) are now included in RNTuple.

Zig-zag encoding is now part of RNTuple but varints aren't.


Resolving patch:
diff --git a/src/uproot/models/RNTuple.py b/src/uproot/models/RNTuple.py
--- a/src/uproot/models/RNTuple.py
+++ b/src/uproot/models/RNTuple.py
@@ -7,6 +7,7 @@
 
 import struct
 import zlib
+from collections import defaultdict
 
 import numpy
 
@@ -50,7 +51,7 @@ def _keys(self):
         keys = []
         field_records = self.header.field_records
         for i, fr in enumerate(field_records):
-            if fr.parent_field_id == i:
+            if fr.parent_field_id == i and fr.type_name != "":
                 keys.append(fr.field_name)
         return keys
 
@@ -144,8 +145,6 @@ def header(self):
             self._header = h
             assert h.crc32 == zlib.crc32(self._header_chunk.raw_data[:-4])
 
-        # cursor = self._header_cursor.copy()
-        # cursor.debug(self._header_chunk)
         return self._header
 
     @property
@@ -238,15 +237,23 @@ def base_col_form(self, cr, col_id, parameters=None):
     def col_form(self, field_id):
         ak = uproot.extras.awkward()
 
-        # FIXME remove this ugly logic
-        rel_crs = []
-        rel_crs_idxs = []
-        for i, cr in enumerate(self.header.column_records):
-            if cr.field_id == field_id:
-                rel_crs.append(cr)
-                rel_crs_idxs.append(i)
-            if cr.field_id > field_id:
-                break
+        if field_id in self._column_recorrds_dict:
+            rel_crs = self._column_recorrds_dict[field_id]["rel_crs"]
+            rel_crs_idxs = self._column_recorrds_dict[field_id]["rel_crs_idxs"]
+        elif field_id in self._alias_columns_dict:
+            rel_crs = self._column_recorrds_dict[self._alias_columns_dict[field_id]][
+                "rel_crs"
+            ]
+            rel_crs_idxs = self._column_recorrds_dict[
+                self._alias_columns_dict[field_id]
+            ]["rel_crs_idxs"]
+        else:
+            raise (
+                RuntimeError(
+                    f"The filed_id: {field_id} is missing both from the columns records and the alias columns."
+                )
+            )
+
         if len(rel_crs) == 1:  # base case
             return self.base_col_form(rel_crs[0], rel_crs_idxs[0])
         elif (
@@ -269,7 +276,7 @@ def field_form(self, this_id, seen):
 
         field_records = self.header.field_records
         this_record = field_records[this_id]
-        seen.append(this_id)
+        seen.add(this_id)
         structural_role = this_record.struct_role
         if (
             structural_role == uproot.const.rntuple_role_leaf
@@ -279,30 +286,23 @@ def field_form(self, this_id, seen):
             # n.b. the split may happen in column
             return self.col_form(this_id)
         elif structural_role == uproot.const.rntuple_role_leaf:
-            # std::array
-            child_id = next(
-                filter(
-                    lambda i: field_records[i].parent_field_id == this_id,
-                    range(this_id + 1, len(field_records)),
-                )
-            )
+            # std::array it only has one child
+            if this_id in self._related_ids:
+                child_id = self._related_ids[this_id][0]
+
             inner = self.field_form(child_id, seen)
             return ak.forms.RegularForm(inner, this_record.repetition)
         elif structural_role == uproot.const.rntuple_role_vector:
-            keyname = self.col_form(this_id)
-            child_id = next(
-                filter(
-                    lambda i: field_records[i].parent_field_id == this_id,
-                    range(this_id + 1, len(field_records)),
-                )
-            )
+            keyname = f"column-{this_id}"
+            #  this only has one child
+            if this_id in self._related_ids:
+                child_id = self._related_ids[this_id][0]
             inner = self.field_form(child_id, seen)
             return ak.forms.ListOffsetForm("u32", inner, form_key=keyname)
         elif structural_role == uproot.const.rntuple_role_struct:
             newids = []
-            for i, fr in enumerate(field_records):
-                if i not in seen and fr.parent_field_id == this_id:
-                    newids.append(i)
+            if this_id in self._related_ids:
+                newids = self._related_ids[this_id]
             # go find N in the rest, N is the # of fields in struct
             recordlist = [self.field_form(i, seen) for i in newids]
             namelist = [field_records[i].field_name for i in newids]
@@ -310,9 +310,8 @@ def field_form(self, this_id, seen):
         elif structural_role == uproot.const.rntuple_role_union:
             keyname = self.col_form(this_id)
             newids = []
-            for i, fr in enumerate(field_records):
-                if i not in seen and fr.parent_field_id == this_id:
-                    newids.append(i)
+            if this_id in self._related_ids:
+                newids = self._related_ids[this_id]
             recordlist = [self.field_form(i, seen) for i in newids]
             return ak.forms.UnionForm("i8", "i64", recordlist, form_key=keyname)
         else:
@@ -325,7 +324,7 @@ def to_akform(self):
         field_records = self.header.field_records
         recordlist = []
         topnames = self.keys()
-        seen = []
+        seen = set()
         for i in range(len(field_records)):
             if i not in seen:
                 recordlist.append(self.field_form(i, seen))
@@ -430,7 +429,7 @@ def read_col_page(self, ncol, cluster_i):
         if zigzag:
             res = from_zigzag(res)
         elif delta:
-            numpy.cumsum(res)
+            res = numpy.cumsum(res)
         return res
 
     def arrays(
@@ -457,17 +456,56 @@ def arrays(
             [c.num_entries for c in clusters[start_cluster_idx:stop_cluster_idx]]
         )
 
+        self._alias_columns_dict = {
+            el.field_id: el.physical_id
+            for i, el in enumerate(self.header.alias_columns)
+        }
+        self._column_recorrds_dict = {}
+        self._column_recorrds_dict = {
+            el.field_id: {
+                "rel_crs": [
+                    *(self._column_recorrds_dict.get(el.field_id) or {}).get(
+                        "rel_crs", []
+                    ),
+                    el,
+                ],
+                "rel_crs_idxs": [
+                    *(self._column_recorrds_dict.get(el.field_id) or {}).get(
+                        "rel_crs_idxs", []
+                    ),
+                    i,
+                ],
+            }
+            for i, el in enumerate(self.header.column_records)
+        }
+
+        self._related_ids = defaultdict(list)
+        for i, el in enumerate(self.header.field_records):
+            if el.parent_field_id != i:
+                self._related_ids[el.parent_field_id].append(i)
+
         form = self.to_akform().select_columns(filter_names)
         # only read columns mentioned in the awkward form
         target_cols = []
         container_dict = {}
         _recursive_find(form, target_cols)
-        for i, cr in enumerate(self.column_records):
-            key = f"column-{i}"
-            dtype_byte = cr.type
-            if key in target_cols:
+        for key in target_cols:
+            if "column" in key:
+                key_nr = int(key.split("-")[1])
+                if key_nr in self._column_recorrds_dict:
+                    id = key_nr
+                elif key_nr in self._alias_columns_dict:
+                    id = self._alias_columns_dict[key_nr]
+                else:
+                    raise (
+                        RuntimeError(
+                            f"The key: {key} is missing both from the columns records and the alias columns."
+                        )
+                    )
+
+                dtype_byte = self._column_recorrds_dict[id]["rel_crs"][0].type
                 content = self.read_col_pages(
-                    i, range(start_cluster_idx, stop_cluster_idx)
+                    id, range(start_cluster_idx, stop_cluster_idx)
                 )
                 if dtype_byte == uproot.const.rntuple_col_type_to_num_dict["switch"]:
                     kindex, tags = _split_switch_bits(content)


Test patch:
diff --git a/tests/test_0962_RNTuple_update.py b/tests/test_0962_RNTuple_update.py
--- a/tests/test_0962_RNTuple_update.py
+++ b/tests/test_0962_RNTuple_update.py
@@ -2,6 +2,7 @@
 
 import pytest
 import uproot
+import awkward as ak
 import skhep_testdata
 import numpy as np
 
@@ -35,3 +36,31 @@ def test_new_support_RNTuple_split_int16_reading():
         assert df.one_integers[-1] == 1
         assert np.all(np.unique(df.one_integers[: len(df.one_integers) // 2]) == [2])
         assert np.all(np.unique(df.one_integers[len(df.one_integers) / 2 + 1 :]) == [1])
+
+
+pytest.importorskip("zstandard")
+
+
+@pytest.mark.xfail(
+    reason="Uproot can now read the data from event files (CMS/Atlas), but this test fails because the column matching logic is wrong.",
+    strict=True,
+)
+def test_new_support_RNTuple_event_data():
+    with uproot.open(
+        "https://xrootd-local.unl.edu:1094//store/user/AGC/nanoaod-rntuple/zstd/TT_TuneCUETP8M1_13TeV-powheg-pythia8/cmsopendata2015_ttbar_19980_PU25nsData2015v1_76X_mcRun2_asymptotic_v12_ext3-v1_00000_0000.root"
+    ) as f:
+        obj = f["Events"]
+        df = obj.arrays(["nTau"])
+        assert len(df) == 1334428
+        assert ak.to_list(df["nTau"][:10]) == [
+            0,
+            0,
+            2,
+            0,
+            1,
+            1,
+            1,
+            1,
+            2,
+            0,
+        ]