feat: implement virtual arrays

[ikrommyd]

This is a PR that aims to implement virtual buffers that are materialized when their actual data is needed. Getting those data would typically be an expensive disk read.
The description shall be updated in the future and become a lot more analytical.

[pfackeldey]

@pfackeldey This should be ready for the first review round. In the meantime, I'll be writing unit tests the following week.

cc @ianna @agoose77 @jpivarski

Let's first figure out unknown_length handling with virtual arrays. Then, we should have all required functionality in place. I'll start a review afterwards! It looks already in very good shape from first sight 🎉

[ikrommyd]

@pfackeldey This should be ready for the first review round. In the meantime, I'll be writing unit tests the following week.
cc @ianna @agoose77 @jpivarski

Let's first figure out unknown_length handling with virtual arrays. Then, we should have all required functionality in place. I'll start a review afterwards! It looks already in very good shape from first sight 🎉

Yup agreed. I wrote this comment before I realised we have this problem.

[pfackeldey]

These changes are great! Thank you very much @ikrommyd 🎉
I have just minor comments that mainly reduce a few lines of code and question in rare cases if a delay of materialization is needed?

[ikrommyd]

Needs unit tests and some analysis-like testing. AGC and coffea benchmarks look good at the moment.

[ikrommyd]

Does anyone know a good place to add a ak.materialize(stuff) before a numba compiled ArrayBuilder operation is run. In arrayview.py perhaps?
For example:

@numba.njit
def find_4lep(events_leptons, builder):
    """Search for valid 4-lepton combinations from an array of events * leptons {charge, ...}

    A valid candidate has two pairs of leptons that each have balanced charge
    Outputs an array of events * candidates {indices 0..3} corresponding to all valid
    permutations of all valid combinations of unique leptons in each event
    (omitting permutations of the pairs)
    """
    for leptons in events_leptons:
        builder.begin_list()
        nlep = len(leptons)
        for i0 in range(nlep):
            for i1 in range(i0 + 1, nlep):
                if leptons[i0].charge + leptons[i1].charge != 0:
                    continue
                for i2 in range(nlep):
                    for i3 in range(i2 + 1, nlep):
                        if len({i0, i1, i2, i3}) < 4:
                            continue
                        if leptons[i2].charge + leptons[i3].charge != 0:
                            continue
                        builder.begin_tuple(4)
                        builder.index(0).integer(i0)
                        builder.index(1).integer(i1)
                        builder.index(2).integer(i2)
                        builder.index(3).integer(i3)
                        builder.end_tuple()
        builder.end_list()

    return builder

This actually materializes and we don't have to materialize ourselves but I don't think that's always gonna happen. In this case it's going to materialize all the buffers of whatever events_leptons is. For external numba compiled functions we can't do anything but for ArrayBuilder stuff I think we could add an intentional ak.materialize call before the kernel runs.

Oh in _lookup.py maybe? A lookup is going to run before calling the arraybuilder right? So we can just materialize the whole layout that goes into the lookup.

[ikrommyd]

Oh I take that back actually, it may be best to do the opposite. Raise in error in the lookup and prompt the user to intentionally materialize.

[pfackeldey]

Oh I take that back actually, it may be best to do the opposite. Raise in error in the lookup and prompt the user to intentionally materialize.

I like that! 👍