Read into an existing numpy array with uproot4

[mpad]

Hello,

I have a workflow where I use uproot3 to read in data from root files directly into existing big in-memory arrays.
I was using interpretation like in :

# declaring this for a bunch of variables : 
mydic[var_name] = tree[var_name].interpretation.toarray(abigarray)
# later :
tree.arrays(branches=mydic)

and as expected all the abigarray are filled in-place.

Is there a way to do that with uproot4 ? I didn't find it in the doc.
I've seen a AsDtypeInPlace class in the interpretation module, but it's currently not implemented so I'm affraid this is not yet possible.
Did I miss something ?

Cheers,

P-A

[jpivarski]

That's right: I just checked and this is indeed unimplemented.

https://github.com/scikit-hep/uproot4/blob/382bb62f47a05f6111d3787e3235f9e56d9e113e/src/uproot/interpretation/numerical.py#L356-L363

It probably got pushed to low priority because it seemed like this feature was hardly ever used in Uproot 3. The way I gauge that is by user comments, questions, and requests for bug-fixes, such as this one. In fact, it's telling that this has been the default Uproot version since last December and this is the first time the point has been raised that it's completely missing.

So, how important is this in your work? Is it memory or speed prohibitive to let Uproot allocate an array, fill it, copy that into the in-place place where you want it, and then let that temporary array go out of scope?

If it's not actually important, then I might just remove the "reminder" stub for AsDtypeInPlace (i.e. remove the promise that it would ever be implemented).

[mpad]

Thanks for your fast answer !
I'm doing DNN training with inputs having O(100)features x O(100M) entries . This does not fit in memory, so I split the inputs in multiple files, read data from 1 file into a big array, perform a maximum of in-place operations on this array (essentially pre-processing) and finally feed the array to the NN library (tensorflow in my case) and then continue on the next file, re-using the array I prepared.
This is a bit complex but I couldn't find a better compromise in terms of performance (reading from file is fast, calculations on big arrays with numpy/numexpr are efficient) and flexibility (typically, I can easily chose to not load all the features I have in the input).
By using as many in-place operations as possible, I can have bigger arrays and thus good perfs.

So in this particular case, yes being able to read directly into a given array is important.

I don't claim I have an optimal solution, but it is based on the idea that when we have to split data into chunks, the bigger the chunks, the better because we can minimize the repetitions of read and preprocessing operations.
To be honest I'm not 100% sure this idea is correct, but if it is I would expect people dealing with large datasets would try to maximize their array and control their memory usage (thus avoid copies and allocation/deallocation of temporary arrays)...
On the other hand, I maybe wrong, maybe some solutions based on multithreading are even better. In this case I'd understand this feature I need is not a priority at all.

Cheers,

P-A

[jpivarski]

There is a limit in flexibility: AsDtypeInPlace would only ever work as a replacement for AsDtype: it wouldn't be possible to in-place fill a jagged array.

Taking as a given that you can't do the DNN training in batches—that the DNN training has to be all in memory at once—it's still not strictly necessary for Uproot to write directly into the DNN array. For instance,

>>> events = uproot.open("Zmumu.root:events")
>>> pretend_this_is_big = np.empty(events.num_entries)
>>> for batch, report in events.iterate(["px1"], step_size=10, report=True):
...     pretend_this_is_big[report.tree_entry_start:report.tree_entry_stop] = batch["px1"]
... 
>>> pretend_this_is_big
array([-41.19528764,  35.11804977,  35.11804977, ...,  32.37749196,
        32.37749196,  32.48539387])

You can preallocate your huge array (called pretend_this_is_big in the above example) and fill it in arbitrarily small batches (10 entries in this example, which is ridiculously small). The memory footprint at any given time is the size of the huge array plus the size of one batch (or maybe two batches because it's not explicitly del'ed at the end of the for loop). We're taking it as a given that the huge array must exist all at once, but the batches can be arbitrarily small. Granted, if you make the batches too small, the filling process will be slow. This is a technique that trades memory for speed.

Is there not a reasonable compromise between memory and speed here? Suppose you need to have a TB in RAM for the DNN's array. Uproot is asymptotically fast for batches about 100 kB or so; let's make it comfortable and say you use batches of 10 or 100 MB. Those batches are still negligible compared to your TB training array. (Still true if we're being more realistic and it's actually 10 GB or so.)

(See TTree.iterate for step_size, which can be set to a size in memory, rather than a number of entries, and report to get the current index of iteration.)

I'm mentioning this because if this works, you can be up and running today. If you have to wait for AsDtypeInPlace to be implemented, you'd have to wait longer. (Unless you want to implement it; I can help if you're so inclined.)

[mpad]

Hello,

Thanks for the detailed suggestion ! I'm just a bit worried that repeating the instruction pretend_this_is_big[atart:stop] = batch[varname] for many variables and steps might kill the performances... but possibly this won't be a big impact and I guess I'll try anyway.

I'm not sure how difficult for me it would be to implement AsDtypeInPlace : if you can give me a few pointers, where to look in the code, I'll get an idea if I can help or not.

[jpivarski]

AsDtypeInPlace inherits from two superclasses, AsDtype and through that, Numerical. There are three steps in interpreting a rectilinear array:

1. Taking the decompressed data from a TBasket and interpreting it as the right dtype. This is implemented in AsDtype and not Numerical, and the implementation is AsDtype.basket_array. This method is already zero-copy, so there would be no reason to override it in AsDtypeInPlace. All it does (apart from calling any user-defined hooks) is data.view(dtype).reshape((-1,) + shape). So you can leave this as an inherited function.
2. Copying each interpreted TBasket into a single output array. For in-place operations, this would be the given array, supplied by a user. The default implementation would be inherited from Numerical (AsDtype does not override it), in Numerical.final_array. This implementation would either need to be overridden or modified to work differently in AsDtypeInPlace because this implementation creates an output = library.empty((length,), self.to_dtype). You would want to derive the output from a user-supplied _to_fill, like this: output = self._to_fill[:length * self.to_dtype.itemsize].view(self.to_dtype). (Ensure that the _to_fill is large enough!) There's a lot of slicing logic here that would be best not to copy-paste, so perhaps the Numerical implementation could check to see if it has hasattr(self, "_to_fill") and if it does, use that, if it doesn't, create a new array, and only the AsDtypeInPlace would have a _tofill?
3. The final_array method calls output = library.finalize(output, branch, self, entry_start, entry_stop), which lets NumPy, Awkward, or Pandas have the last word on what the output looks like. AsDtypeInPlace can only work with library="np" (an exception should be raised early if any other library is used), and NumPy.finalize seems to let non-jagged arrays through without modification, so this should be good as long as you enforce that AsDtypeInPlace can't be used with anything but library="np".

With these modifications, the following data transformations would happen, from start to finish:

• Data read from disk (if you're using the default file_handler=uproot.MemmapSource) is actually a page of virtual memory in the operating system (the memory-map). That's a minimal-copy technique. The alternative (uproot.MultithreadedFileSource, a plain file handle) is to copy from virtual memory into an array in the Python process.
• If your ROOT file has any compression, it must be copied into a new buffer to decompress it, one TBasket at a time. There's no in-place alternative: decompressed data can have a different size (hopefully bigger!) than the original, so it must be a new array. If your ROOT file has no compression, the "uncompressed TBaskets" might still be the memory-map.
• By inheriting AsDtype.basket_array, you just view those bytes in the correct interpretation, which for ROOT is always big-endian.
• By overriding or otherwise modifying the behavior of Numerical.final_array, you copy the interpreted TBasket arrays into the output array, and the bytes are flipped to native-endian during the copy. There is no way around doing this copy—even if the ROOT file were uncompressed—since the TBaskets are not concatenated, but you need the output to be concatenated. The fact that we get to do the endianness in this copy is a nice bonus. (The endianness conversion is in the distinction between from_dtype and to_dtype: the first is big-endian, the second is native-endian.)

So there is copying happening—from the uncompressed TBasket arrays (which might or might not be OS pages) to the desired output array. But the TBasket arrays are only as large as the baskets (rarely as large as a MB each), and there is only one big array. It's like the "many small temporary arrays filling parts of the big array and then being deleted" approach I proposed above, except with one fewer (small) copy, from TBaskets to batches. With AsDtypeInPlace, data are copied directly from the TBaskets into your given array.

The AsDtypeInPlace class will need to take an array in its __init__ instead of a to_dtype, and it might infer the from_dtype as the big-endian version of the given array's dtype. Instead of having _to_type as an instance attribute, it would have _to_fill as the instance attribute and would derive to_dtype as a property.

That's all I can think of!

[mpad]

Hi,

With a bit of delay, I could give a try and had something working.
I've opened the MR #487 to discuss if what I did can be a solution

[mpad]

Thanks for helping with this PR and accepting it !
Any idea when a next release including it will be available ?

[jpivarski]

I started the process toward 4.1.7, since I wanted to get a restriction on Awkward version out, too.