Skimming NanoEvents

[nsmith-]

An often-requested feature is a way to skim (filter events and columns) a NanoEvents object. Of course awkward can write parquet files for any awkward array with ak.to_parquet, but often users want to skim into a ROOT file. To do so, the following snip can work, but in the long run we will build a tool into coffea:

def is_rootcompat(a):
    """Is it a flat or 1-d jagged array?"""
    t = ak.type(a)
    if isinstance(t, ak._ext.ArrayType):
        if isinstance(t.type, ak._ext.PrimitiveType):
            return True
        if isinstance(t.type, ak._ext.ListType) and isinstance(t.type.type, ak._ext.PrimitiveType):
            return True
    return False


def uproot_writeable(events):
    """Restrict to columns that uproot can write compactly"""
    out = {}
    for bname in events.fields:
        if events[bname].fields:
            out[bname] = ak.zip(
                {
                    n: ak.packed(ak.without_parameters(events[bname][n]))
                    for n in events[bname].fields
                    if is_rootcompat(events[bname][n])
                }
            )
        else:
            out[bname] = ak.packed(ak.without_parameters(events[bname]))
    return out

An example use is

# https://raw.githubusercontent.com/CoffeaTeam/coffea/master/tests/samples/nano_dy.root
filename = "tests/samples/nano_dy.root"
events = NanoEventsFactory.from_root(
    filename,
    metadata={"dataset": "SomeDataset"},
).events()

events = events[ak.sum(events.Electron.pt > 10, axis=1) > 0]
with uproot.recreate("events.root") as fout:
    fout["Events"] = uproot_writeable(events)

[ikrommyd]

Just commenting a working version for awkward 2 and eager coffea.

import awkward as ak
import uproot
from coffea.nanoevents import NanoEventsFactory

def is_rootcompat(a):
    """Is it a flat or 1-d jagged array?"""
    t = ak.type(a)
    if isinstance(t, ak.types.ArrayType):
        if isinstance(t.content, ak.types.NumpyType):
            return True
        if isinstance(t.content, ak.types.ListType) and isinstance(t.content.content, ak.types.NumpyType):
            return True
    return False


def uproot_writeable(events):
    """Restrict to columns that uproot can write compactly"""
    out = {}
    for bname in events.fields:
        if events[bname].fields:
            out[bname] = ak.zip({n: ak.to_packed(ak.without_parameters(events[bname][n])) for n in events[bname].fields if is_rootcompat(events[bname][n])})
        else:
            out[bname] = ak.to_packed(ak.without_parameters(events[bname]))
    return out


filename = "https://raw.githubusercontent.com/CoffeaTeam/coffea/master/tests/samples/nano_dy.root"
events = NanoEventsFactory.from_root({filename: "Events"}, delayed=False).events()

events = events[ak.sum(events.Electron.pt > 10, axis=1) > 0]
with uproot.recreate("skimmedevents.root") as fout:
    fout["Events"] = uproot_writeable(events)