Example 1: Streamer method is overridden

Example 1: `Streamer` method is overridden#

Step 1: Check the binary data#

Before implementing the reader and factory, we should check the binary data of TOverrideStreamer to understand how the data is stored in the ROOT file:

>>> import uproot
>>> import uproot_custom as uc
>>>
>>> br = uproot.open("demo_data.root")["my_tree:override_streamer"]
>>> bin_arr = br.array(interpretation=uc.AsBinary())
>>> evt0 = bin_arr[0].to_numpy()
>>> evt0
array([  0,   1,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
         0,  18,  52,  86, 120,   0,   0,   0,   0,   0,   0,   0,   0],
      dtype=uint8)

Referring to the Streamer method above, we can see that the binary data contains:

TObject content (10 bytes, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0)
m_int (4 bytes, 0, 0, 0, 0, which is 0)
mask (4 bytes, 18, 52, 86, 120, which is 0x12345678)
m_double (8 bytes, 0, 0, 0, 0, 0, 0, 0, 0, which is 0.0)

These bytes are the data your reader needs to read.

Step 2: Implement C++ `reader` to read binary data#

We can implement a reader named OverrideStreamerReader to do this:

#include <cstdint>
#include <memory>
#include <vector>

#include "uproot-custom/uproot-custom.hh"

using namespace uproot;

class OverrideStreamerReader : public IElementReader {
  public:
    OverrideStreamerReader( std::string name )
        : IElementReader( name )
        , m_data_ints( std::make_shared<std::vector<int>>() )
        , m_data_doubles( std::make_shared<std::vector<double>>() ) {}

    void read( BinaryBuffer& buffer ) {
        // Skip TObject header
        buffer.skip_TObject();

        // Read integer value
        m_data_ints->push_back( buffer.read<int>() );

        // Read a custom added mask value
        auto mask = buffer.read<uint32_t>();
        if ( mask != 0x12345678 )
        {
            throw std::runtime_error( "Error: Unexpected mask value: " +
                                      std::to_string( mask ) );
        }

        // Read double value
        m_data_doubles->push_back( buffer.read<double>() );
    }

    py::object data() const {
        auto int_array    = make_array( m_data_ints );
        auto double_array = make_array( m_data_doubles );
        return py::make_tuple( int_array, double_array );
    }

  private:
    const std::string m_name;
    std::shared_ptr<std::vector<int>> m_data_ints;
    std::shared_ptr<std::vector<double>> m_data_doubles;
};

// Declare the reader
PYBIND11_MODULE( my_reader_cpp, m ) {
    declare_reader<OverrideStreamerReader, std::string>( m, "OverrideStreamerReader" );
}

In the read method, we skip the TObject header, then read the member variables and the mask according to the logic in the Streamer method.

In the data method, we return a py::tuple containing 2 numpy arrays: one for m_int and the other for m_double.
Finally we declare the reader in the PYBIND11_MODULE, so that it can be used in Python.

Step 3: Implement Python `factory`#

To use our OverrideStreamerReader and reconstruct the final awkward array, we need to implement a corresponding factory. We can implement a factory named OverrideStreamerFactory to do this.

A factory requires at least 3 methods: build_factory, build_cpp_reader and make_awkward_content. An optional method make_awkward_form can be implemented to enable dask functionality.

First, import necessary modules:

import awkward.contents
import awkward.forms
from uproot_custom import Factory
from .my_reader_cpp import OverrideStreamerReader

The my_reader_cpp is the compiled C++ module containing our OverrideStreamerReader.

Implement `build_factory`#

We can make an assumpion that the fName of the TStreamerInfo is TOverrideStreamer for our class. If the fName matches, we return a tree config dictionary containing the factory and the name of the corresponding reader. Otherwise, we return None to let other factories have a chance to handle current class.

class OverrideStreamerFactory(Factory):
    @classmethod
    def build_factory(
        cls,
        top_type_name: str,
        cur_streamer_info: dict,
        all_streamer_info: dict,
        item_path: str,
        **kwargs,
    ):
        fName = cur_streamer_info["fName"]
        if fName != "TOverrideStreamer":
            return None

        return cls(fName) # Factory takes `name: str` as constructor argument

Tip

In production, you may want use item_path to make a more accurate identification of whether the current class is the one you want to handle:

class OverrideStreamerFactory(Factory):
    @classmethod
    def build_factory(
        cls,
        top_type_name: str,
        cur_streamer_info: dict,
        all_streamer_info: dict,
        item_path: str,
        **kwargs,
    ):
        if item_path != "/my_tree:override_streamer":
            return None

        return cls(fName)

Implement `build_cpp_reader`#

Implement build_cpp_reader to create an instance of OverrideStreamerReader:

def build_cpp_reader(self):
    return OverrideStreamerReader(self.name)

Implement `make_awkward_content`#

Implement make_awkward_content to construct awkward contents from the raw data returned by the reader:

def make_awkward_content(self, raw_data):
    int_array, double_array = raw_data

    return awkward.contents.RecordArray(
        [
            awkward.contents.NumpyArray(int_array),
            awkward.contents.NumpyArray(double_array),
        ],
        ["m_int", "m_double"],
    )

The raw_data is the object returned by the data method of the reader. In our example, it is a py::tuple containing 2 numpy arrays, as illustrated above.

(Optional) Implement `make_awkward_form`#

The make_awkward_form method is optional, but it is easy to implement, since the awkward.forms is similar to awkward.contents:

def make_awkward_form(self):
    return awkward.forms.RecordForm(
        [
            awkward.forms.NumpyForm("int32"),
            awkward.forms.NumpyForm("float64"),
        ],
        ["m_int", "m_double"],
    )

Step 4: Register target branch and the `factory`#

Finally, we need to register the branch we want to read with uproot-custom, and also register the OverrideStreamerFactory so that it can be used by uproot-custom.

We can do this by adding the following code in the __init__.py of your package:

from uproot_custom import registered_factories, AsCustom

AsCustom.target_branches |= {
    "/my_tree:override_streamer",
}

registered_factories.add(OverrideStreamerFactory)

Step 5: Read data with `uproot`#

Now we can read the data using uproot as usual:

>>> b = uproot.open("demo_data.root")["my_tree:override_streamer"]
>>> arr = b.array()

Example 1: Streamer method is overridden

Contents

Example 1: Streamer method is overridden#

Step 1: Check the binary data#

Step 2: Implement C++ reader to read binary data#

Step 3: Implement Python factory#

Implement build_factory#

Implement build_cpp_reader#

Implement make_awkward_content#

(Optional) Implement make_awkward_form#

Step 4: Register target branch and the factory#

Step 5: Read data with uproot#