Incorrect writing of jagged arrays into a TTree

[vpec0]

When writing a jagged array into a TTree in an output ROOT file, if the array is larger than the TBasket, segmentation violation occurs when attempting to read the data in a ROOT session. UpRoot stores incorrect number of baskets used which possibly leads to the issue in ROOT.

Simple example:

of=ur.recreate('test.root')
of['tree']={'wfm':ak.from_regular(ak.Array(np.arange(32000*5*6,dtype='float').reshape(32000,5,6)),axis=-2)}
of['tree'].num_baskets # this is equal to 1
of.close()

After opening in ROOT (note the size of the basket and number of baskets for branch wfm vs the size of the branch):

$ root test.root
root [0] 
Attaching file test.root as _file0...
(TFile *) 0x7f9e8df4a730
root [1] tree->Print()
******************************************************************************
*Tree    :tree      :                                                        *
*Entries :    32000 : Total =         7937482 bytes  File  Size =    1410225 *
*        :          : Tree compression factor =   5.63                       *
******************************************************************************
*Br    0 :nwfm      : nwfm/I                                                 *
*Entries :    32000 : Total  Size=     128552 bytes  File Size  =        729 *
*Baskets :        1 : Basket Size=      32000 bytes  Compression= 175.68     *
*............................................................................*
*Br    1 :wfm       : wfm[nwfm][6]/D                                         *
*Entries :    32000 : Total  Size=    7808656 bytes  File Size  =    1408155 *
*Baskets :        1 : Basket Size=      32000 bytes  Compression=   5.54     *
*............................................................................*
root [2] tree->Draw("wfm")
Info in <TCanvas::MakeDefCanvas>:  created default TCanvas with name c1
Warning in <TBasket::ReadBasketBuffers>: basket:wfm has fNevBuf=32000 but fEntryOffset=0, pos=22484, len=7808078, fNbytes=1408155, fObjlen=7808008, trying to repair

 *** Break *** segmentation violation
...

I note here that if the stored array is not jagged, ROOT gives no segmentation violation even though the size of the branch is larger than a single basket and the number of baskets for the branch is still 1.

[vpec0]

To add here, the data cannot be retrieved even with UpRoot.

>>> with ur.open('test.root:tree') as t :
...     d = t.arrays()
...
Traceback (most recent call last):
  File "/usr/local/lib/python3.13/site-packages/uproot/interpretation/numerical.py", line 359, in basket_array
    output = data.view(dtype).reshape((-1, *shape))
ValueError: cannot reshape array of size 20000 into shape (6)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/lib/python3.13/site-packages/uproot/behaviors/TBranch.py", line 889, in arrays
    _ranges_or_baskets_to_arrays(
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        self,
        ^^^^^
    ...<9 lines>...
        interp_options,
        ^^^^^^^^^^^^^^^
    )
    ^
  File "/usr/local/lib/python3.13/site-packages/uproot/behaviors/TBranch.py", line 3204, in _ranges_or_baskets_to_arrays
    uproot.source.futures.delayed_raise(*obj)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
  File "/usr/local/lib/python3.13/site-packages/uproot/source/futures.py", line 38, in delayed_raise
    raise exception_value.with_traceback(traceback)
  File "/usr/local/lib/python3.13/site-packages/uproot/behaviors/TBranch.py", line 3140, in basket_to_array
    basket_array = interpretation.basket_array(
        basket.data,
    ...<6 lines>...
        interp_options,
    )
  File "/usr/local/lib/python3.13/site-packages/uproot/interpretation/jagged.py", line 180, in basket_array
    content = self._content.basket_array(
        data, None, basket, branch, context, cursor_offset, library, options
    )
  File "/usr/local/lib/python3.13/site-packages/uproot/interpretation/numerical.py", line 361, in basket_array
                raise ValueError(
    ...<3 lines>...
                ) from err
ValueError: basket 0 in tree/branch /tree;1:wfm has the wrong number of bytes (160000) for interpretation AsDtype("('>f8', (6,))")
in file test.root

[ariostas]

Thank you for reporting this, @vpec0! I'll look into it.

[ariostas]

As far as I know, Uproot has no mechanism for splitting an array into multiple baskets, so it makes sense that it always writes a single basket for the branch.

It doesn't seem like the issue is related to the size of the array being larger than what the TBasket can hold, because even a small array fails to be written correctly.

>>> of=ur.recreate('test.root')
... of['tree']={'wfm': ak.from_regular(ak.Array(np.arange(32*5*6,dtype='float').reshape(-1,5,6)))}
... print(of['tree'].num_baskets) # this is equal to 1
... of.close()
... with ur.open('test.root:tree') as t:
...     d = t.arrays()

So it looks like the issue happens when it is a jagged array of regular arrays. I'm surprised that this wasn't tested.

I'm not very familiar with the TTree writing parts, so it might take me some time to figure it out.

In the meantime, I'd invite you to switch to using RNTuples instead of a TTree. We intend to fully support RNTuples, as opposed to the limited support that we offer for TTrees. Here's how you can do it.

>>> of=ur.recreate('test.root')
... data = {'wfm': ak.from_regular(ak.Array(np.arange(32*5*6,dtype='float').reshape(-1,5,6)))}
... of.mkrntuple("rntuple", data)
... of.close()
... with ur.open('test.root:rntuple') as t:
...     d = t.arrays()

[ikrommyd]

As far as I know, Uproot has no mechanism for splitting an array into multiple baskets, so it makes sense that it always writes a single basket for the branch.

Btw...there have been cases in analysis code where having this mechanism in uproot would have been very beneficial instead of looping over chunks of the array and appending to the tree to write multiple tbaskets manually.

[vpec0]

In the meantime, I'd invite you to switch to using RNTuples instead of a TTree. We intend to fully support RNTuples, as opposed to the limited support that we offer for TTrees.

I don't mind using RNTuples, however, in my analysis, I do have nested records -- example:

print(data.typestr)
23 * {run: uint32, event: uint32, hasT0: uint16, end_xyz: 3 * float32, t0: float32, flash_time: var * float32, avg_wfm: var * float64, peak: var * {idx: int64, prominence: float64, ipr: float64}, wfm: var * 1024 * float64}

mkrntuple cannot write that data for me:

...
  File "/usr/local/lib/python3.13/site-packages/uproot/writing/_cascadentuple.py", line 496, in _build_field_col_records
    raise NotImplementedError(f"Form type {type(akform)} cannot be written yet")
NotImplementedError: Form type <class 'awkward.forms.indexedform.IndexedForm'> cannot be written yet

Or is there a workaround?

[ariostas]

@ikrommyd yeah, that's definitely a useful feature and I'm planning to add it for RNTuples.

@vpec0 Sorry about that, it currently supports nested records, but I'm still missing a few Awkward form types. The one you point out is pretty easy to implement, so I'll aim to add it by the next release.

[ariostas]

@vpec0 I opened #1493 to add support for IndexedArrays. We'll have a release of Uproot next week, so it should work by then.

[ianna]

Hi @vpec0 Could you please check the new release and let us know if it is fixed in #1493? Thanks!

[vpec0]

I have tested the workaround that uses RNTuple with the latest uproot version, 5.6.5. I still do get an error when attempting to write the data:

    print(data)
    outf.mkrntuple(tree_name,data)

gives the following error:

    outf.mkrntuple(tree_name,data)
    ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.13/site-packages/uproot/writing/writable.py", line 1376, in mkrntuple
    ntuple.extend(ak_data)
    ~~~~~~~~~~~~~^^^^^^^^^
  File "/usr/local/lib/python3.13/site-packages/uproot/writing/writable.py", line 2180, in extend
    self._cascading.extend(self._file, self._file.sink, data)
    ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.13/site-packages/uproot/writing/_cascadentuple.py", line 971, in extend
    content = data_buffers[f"{next_barekey}-data"]
              ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'node8-data'
Command exited with non-zero status 1

@ianna I am assuming nothing has been done on the TTree side, so no point to test that, right?

[ariostas]

@vpec0 nothing has been done on the TTree side yet.

Would you be able to pickle the data you're trying to save and attach it here so I can debug it? I must have missed something while implementing IndexArrays.

[ariostas]

Nevermind, I know what's wrong. I'll work on fixing it.

[vpec0]

Good, because I could not reproduce the error after I had pickled the array.

[ariostas]

@vpec0 could you try using #1496 to see if it finally works?