Numba recursion, ArrayBuilder performance, and writing functions that descend to an an arbitrary list depth

[tamasgal]

This issue arose in #572

I have difficulties with a recursive function using awkward.Arrays. The following example simply traverses an array with an ndim>=2 and produces an array with ndim-1 which contains a boolean entry for each sub-array which starts and ends with the given start and end integers. It works fine without njit but fails in the Numba context as seen below.

def mask_startend(arr, start, end):
    builder = ak.ArrayBuilder()
    _mask_startend(arr, start, end, builder)
    return builder.snapshot()

@nb.njit  # raises 
def _mask_startend(arr, start, end, builder):
    if arr.ndim == 2:  # recursion stop
        for el in arr:
            if len(el) > 0 and el[0] == start and el[-1] == end:
                builder.boolean(True)
            else:
                builder.boolean(False)
        return
    
    for subarray in arr:
        builder.begin_list()
        _mask_startend(subarray, start, end, builder)
        builder.end_list()

This gives without @nb.njit (only [1, 2, 3] and [1, 3] start with 1 and 3):

>>> mask_startend(ak.Array([[[1,2,3], [1,2]], [[1, 3]], [[1, 2], [3, 4]]]), 1, 3).tolist()
[[True, False], [True], [False, False]]

But gives an error when JITted. It seems that the eq function from el[0] == start is failing due to type inference. el[0] should report its type as a simple int64 but somehow it's that array view which of course doesn't have an == implementation.
I guess that this is coming from type inference confusion caused by the iterator interface, but I am not sure.

Any hints or ideas are appreciated!

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function eq>) found for signature:

 >>> eq(ak.ArrayView(ak.NumpyArrayType(array(int64, 1d, A), none, {}), None, ()), int64)

There are 32 candidate implementations:
      - Of which 30 did not match due to:
      Overload of function 'eq': File: <numerous>: Line N/A.
        With argument(s): '(ak.ArrayView(ak.NumpyArrayType(array(int64, 1d, A), none, {}), None, ()), int64)':
       No match.
      - Of which 2 did not match due to:
      Operator Overload in function 'eq': File: unknown: Line unknown.
        With argument(s): '(ak.ArrayView(ak.NumpyArrayType(array(int64, 1d, A), none, {}), None, ()), int64)':
       No match for registered cases:
        * (bool, bool) -> bool
        * (int8, int8) -> bool
        * (int16, int16) -> bool
        * (int32, int32) -> bool
        * (int64, int64) -> bool
        * (uint8, uint8) -> bool
        * (uint16, uint16) -> bool
        * (uint32, uint32) -> bool
        * (uint64, uint64) -> bool
        * (float32, float32) -> bool
        * (float64, float64) -> bool
        * (complex64, complex64) -> bool
        * (complex128, complex128) -> bool

During: typing of intrinsic-call at <ipython-input-8-00220a4074bc> (10)

File "<ipython-input-8-00220a4074bc>", line 10:
def _mask_startend(arr, start, end, builder):
    <source elided>
        for el in arr:
            if len(el) > 0 and el[0] == start and el[-1] == end:
            ^

[jpivarski]

Without knowing your specific problem, I put a note about recursion in the other issue: #572 (comment)

I was actually in a Numba developer's meeting, but we ran overtime and I didn't get to ask about recursion.

Also, because I hadn't investigated it before, I checked to see if there's any documentation on it. There is:

https://numba.pydata.org/numba-doc/latest/proposals/typing_recursion.html

The last note on that page about "limitations" seems to be about what I found in the code (see above comment) about the call stack.

I'm going to see if I can cast your problem in the form that has been implemented. Failing that, I'll close this issue because it's actually a Numba issue.

[tamasgal]

Yes thanks! Saw it already :)

Btw. I also found the NBEP 6, but could not figure out how to use type annotation with complex structures like these awkward arrays. If you have any examples for that, that would be very helpful! I already worked with type annotated Numba-JIT functions before, but only with primitive types and NumPy arrays.

[jpivarski]

The first quick thing is that if array.ndim == 2, then

el[0] == start

would be an array comparison, which would return an array of booleans (is the first element of el[0] equal to start? is the second? is the third?), which can't be used as a predicate, anyway. You don't see that error message because "np.equal(Awkward Array, number)" hasn't been implemented. That's what I was asking about in the meeting: I want to implement all of these by taking advantage of Numba's existing ufuncs (such as np.equal) and I was getting hints on how to do that.

So as a first step, you either want something different from

el[0] == start and el[-1] == end

or perhaps you instead wanted this case to be ndim == 1:

if arr.ndim == 1:  # recursion stop

[tamasgal]

Ah, that's actually a doubly nested for-loop (which is in fact the lowest dimension to be supported), so el is a sub-array. The variable names are a bit misleading, sorry 🙈. So el is a subarray and then I am looking for the first and last element if those are equal to start and end.

So if I understood correctly, the problem is that if a higher dimensional (arr.ndim > 2) is passed in, the if arr.ndim == 2 is not skipped because there is no logic resolution of the code at that stage, so it goes further and still tries to find len(el) > 0 which is then in fact a higher dimensional one (and actually doesn't make sense in the algorithm) and does not implement equality (we wouldn't need them anyways but it still tries to).
The main problem therefore is that that specific part of the code is only valid for arr.ndim == 2 but cannot be hidden for higher dimns.

    if arr.ndim == 2:  # recursion stop
        for el in arr:
            if len(el) > 0 and el[0] == start and el[-1] == end:
                builder.boolean(True)
            else:
                builder.boolean(False)
        return

[jpivarski]

I should still think about the Numba recursion problem, but you know, this can be solved entirely by Awkward Array primitives, too.

Starting with an example of integer endpoints like

>>> import awkward as ak
>>> example = ak.Array([[[0, 1, 2], []], [[3, 4]], [], [[5], [6, 7, 8, 9]]])
>>> example
<Array [[[0, 1, 2], []], ... 5], [6, 7, 8, 9]]] type='4 * var * var * int64'>

Some of these are empty, so part of the solution will involve avoiding accessing the first or last element of an empty array.

>>> ak.num(example, axis=-1)
<Array [[3, 0], [2], [], [1, 4]] type='4 * var * int64'>
>>> nonempty = ak.num(example, axis=-1) > 0
>>> nonempty
<Array [[True, False], ... [True, True]] type='4 * var * bool'>

If we used this as a cut, it would reduce the length of the inner lists, which would make the final result hard to apply to arrays of the original length. So instead of a standard boolean slice, we'll use a mask-slice:

>>> example.mask[nonempty]
<Array [[[0, 1, 2], None], ... [6, 7, 8, 9]]] type='4 * var * option[var * int64]'>

We want to pick out the first and last elements of each remaining list at the deepest level, whatever that is. As in NumPy, we can use ellipsis (...):

>>> example.mask[nonempty][..., 0]
<Array [[0, None], [3], [], [5, 6]] type='4 * var * ?int64'>
>>> example.mask[nonempty][..., -1]
<Array [[2, None], [4], [], [5, 9]] type='4 * var * ?int64'>

For a given start and stop, like 6 and 9,

>>> start, stop = 6, 9
>>> example.mask[nonempty][..., 0] == start
<Array [[False, None], ... [False, True]] type='4 * var * ?bool'>
>>> example.mask[nonempty][..., -1] == stop
<Array [[False, None], ... [False, True]] type='4 * var * ?bool'>

We need both of those to be true:

>>> good = ((example.mask[nonempty][..., 0] == start) &
...         (example.mask[nonempty][..., -1] == stop))
>>> good
<Array [[False, None], ... [False, True]] type='4 * var * ?bool'>

We still have None values in the output. We could use this as a slice, but it's probably not what you want:

>>> example[good]
<Array [[None], [], [], [[6, 7, 8, 9]]] type='4 * var * option[var * int64]'>

We don't want to pick out the empty arrays as None; we want them to go away, which is what False would do. So just replace None → False:

>>> ak.fill_none(good, False)
<Array [[False, False], ... [False, True]] type='4 * var * bool'>

Putting it all together,

>>> def mask_startstop(arr, start, stop):
...     nonempty = ak.num(arr, axis=-1) > 0
...     good = ((arr.mask[nonempty][..., 0] == start) &
...             (arr.mask[nonempty][..., -1] == stop))
...     return ak.fill_none(good, False)
...
>>> mask_startstop(example, 0, 2)
<Array [[True, False], ... [False, False]] type='4 * var * bool'>
>>> mask_startstop(example, 5, 5)
<Array [[False, False], ... [True, False]] type='4 * var * bool'>

But I'll think about the Numba thing, too, because I want to understand what Numba can and can't do, and why.

[tamasgal]

Thanks for the nice example, that's also very helpful. I already had a somewhat similar solution but wanted to squeeze out the last bits of performance since we need to apply this to very large datasets and I wanted to avoid allocation as much as possible. I will review your workflow and compare it with the current implementation, which contains just hardcoded functions for a few dimensions.

[jpivarski]

Keep in mind that ak.ArrayBuilder is not super-fast. It's compiled, but dynamically typed. Someday, we'll need to have a "TypedArrayBuilder" that takes a type as a constructor argument to build array structures without having to discover the type at runtime.

[tamasgal]

Oh OK, I did not know that. I'll definitely compare the performance with your approach anyways and report back!

[jpivarski]

I struggled with Numba recursion for a while and decided that I don't understand their call-stack restriction. But anyway, you're not going to get the most speed from that because of ArrayBuilder. The absolutely fastest way to do this is to unwrap the array until you're down to just the part that the Numba-compiled function applies to, then wrap the result.

I have an open issue (#516) to make a public interface that would let you write fully general functions like this. Perhaps I should prioritize that. Meanwhile, here's what it would look like, fully written out. The ak.to_layout(x) and x.layout parts extract a layout from an array (the latter only works for ak.Array) and the ak.Array constructors reverse that process. The .purelist_depth is the low-level layout name for ndim. (It's hard to keep names consistent: we're only careful about naming conventions at the high-level ak.Array interface.)

import awkward as ak
import numpy as np
import numba as nb

@nb.njit
def base_case(arr, start, stop):
    out = np.empty(len(arr), np.bool_)
    for i, el in enumerate(arr):
        out[i] = len(el) > 0 and el[0] == start and el[-1] == stop
    return out

def mask_startstop(arr, start, stop):
    def recurse(layout):
        if layout.purelist_depth == 2:
            np_array = base_case(ak.Array(layout), start, stop)
            return ak.layout.NumpyArray(np_array)

        elif isinstance(layout, (
            ak.layout.ListArray32,
            ak.layout.ListArrayU32,
            ak.layout.ListArray64,
        )):
            content = recurse(layout.content)
            return type(layout)(layout.starts, layout.stops, content)

        elif isinstance(layout, (
            ak.layout.ListOffsetArray32,
            ak.layout.ListOffsetArrayU32,
            ak.layout.ListOffsetArray64,
        )):
            content = recurse(layout.content)
            return type(layout)(layout.offsets, content)

        elif isinstance(layout, ak.layout.RegularArray):
            content = recurse(layout.content)
            return ak.layout.RegularArray(content, layout.size)

        else:
            raise NotImplementedError(repr(arr))

    layout = ak.to_layout(arr, allow_record=True, allow_other=False)
    return ak.Array(recurse(layout))

What #516 would do is replace the if statements identifying layout node types with a functional callback, like the one I use internally, but with a guaranteed public interface. (One of the tasks I set for myself today is to change that callback—my internal function is not a stable API, though one should be built on top of it.)

This will be the fastest implementation because

• it does not recreate the offsets or starts/stops at all levels above ndim == 2 (I was wrong in my earlier comment: the base case is ndim == 2); it reuses the old ones—no allocation and no copying.
• it also doesn't create intermediate arrays in the len(el) > 0 and el[0] == start and el[-1] == stop, which my pure high-level Awkward solution did.

It will generally be true that unwrapping/reuse + Numba for the fixed-type things that return NumPy arrays will be the fastest for this combination of reasons, so there isn't really a reason to get Numba recursion to work.

[tamasgal]

Here is a quick comparison, it seems that the Numba version (mask) is still much faster than the awkward whole-meal appraoch (mask_alt). The f.events.tracks.rec_stages is coming from uproot4, I can dig out the sample file from skhep-testdata later!

>>> arr2d = f.events.tracks.rec_stages[0]

>>> arr3d = f.events.tracks.rec_stages

>>> arr2d
<Array [[1, 3, 5, 4], [1, 3, ... [1], [1], [1]] type='56 * var * int64'>

>>> arr3d
<Array [[[1, 3, 5, 4], [1, ... 1], [1], [1]]] type='10 * var * var * int64'>

>>> %timeit km3io.tools.mask(arr2d, startend=(1, 3))
155 µs ± 3.45 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit km3io.tools.mask(arr3d, startend=(1, 3))
197 µs ± 18.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

>>> %timeit km3io.tools.mask_alt(arr2d, 1, 3)
3.54 ms ± 65.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

>>> %timeit km3io.tools.mask_alt(arr3d, 1, 3)
7.04 ms ± 2.54 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

[jpivarski]

I think these samples are too small to be conclusive as performance tests. There are two things that could be quantified: the time spent rearranging metadata (like those layout nodes in my previous comment) and the time spent iterating over arrays. If you do a test with a small dataset over and over (100‒10000 times, above), you're mostly measuring the time spent rearranging metadata. In real applications, that mask_alt function is only going to be called O(1) times—like once, or a hundred times, or something small like that—an anyone can wait 100 µs × a hundred times or 7.04 ms × a hundred times.

Therefore, you want your performance test to run over a large dataset a small number of times. The first number out of %timeit should be at least seconds.

[tamasgal]

Yes you are of course right! However, this will be part of an iterative pipeline where each "event" is read-in one-by-one or chunk-wise (depending on the available memory and the type of the processing chain). I will talk to the users to find a realistic use-case and do the performance checks on that. I also need to study your comment at #580 (comment) which I overlooked since it was posted the same time I answered ;)

[tamasgal]

Here is a quick comparison of mask and mask_alt. To quickly recap, the mask uses Numba and is currently hardcoded to only work with 2 and 3 dimensions (btw. the actual mask function we use implements also other condition checking mechanisms, this is a simplified version which only implements the startend-check and is compared to the awkward-slicing-magic mask_alt approach provided by Jim):

def mask(arr, start, end):
    builder = ak.ArrayBuilder()
    # Numba has limited recursion support, so this is hardcoded
    if arr.ndim == 2:
        _mask2d(arr, builder, start, end)
    else:
        _mask3d(arr, builder, start, end)
    return builder.snapshot()

@nb.njit
def _mask3d(arr, builder, start, end):
    for subarray in arr:
        builder.begin_list()
        _mask2d(subarray, builder, start, end)
        builder.end_list()

@nb.njit
def _mask2d(arr, builder, start, end):
    for els in arr:
        if len(els) > 0 and els[0] == start and els[-1] == end:
            builder.boolean(True)
        else:
            builder.boolean(False)

And the mask_alt using awkward-only high-level functions:

def mask_alt(arr, start, end):
    nonempty = ak.num(arr, axis=-1) > 0
    mask = ((arr.mask[nonempty][..., 0] == start) & (arr.mask[nonempty][..., -1] == end))
    return ak.fill_none(mask, False)

The benchmark is dead simple:

from collections import defaultdict
import awkward as ak
import numpy as np
import uproot
import matplotlib.pyplot as plt
from skhep_testdata import data_path

# This is a large file with 143381 events, each containing a 128 element
# array which consists of up ~10 elements of `int64`
# An example file with 10 events is available via
# f = uproot.open(data_path("uproot-issue431b.root"))
f = uproot.open("datav5.40.jorcarec.aanet.00006614.root")
arr = f["E/Evt/trks/trks.rec_stages"].array()  # takes 3-4 minutes and 8GB mem to load

ns = 2 ** np.arange(18)
benchmarks = defaultdict(list)
for n in ns:
    print(f"{n}")
    bm_numba = %timeit -o km3io.tools.mask(arr[0:n], startend=(1, 3))
    bm_awkward = %timeit -o km3io.tools.mask_alt(arr[0:n], 1, 3)
    benchmarks["numba"].append(bm_numba)
    benchmarks["awkward"].append(bm_awkward)

fig, ax = plt.subplots()
for lib, benchmark in benchmarks.items():
    ax.scatter(ns, [b.average for b in benchmark], label=lib)
ax.legend();
ax.set_xlabel("number of arrays")
ax.set_ylabel("average execution time / us")
ax.grid()

The overhead seems to be negligible and Numba seems to be ~3 times faster:

[tamasgal]

Yes indeed, factor 3 is quite impressive (and it works for arbitrary dimensions)! "size of arrays" is more descriptive and correct, I agree.

The two cases studied here don't include the one I wrote in #580 (comment), which ought to be the fastest of the three.

Ah, let me add that quickly...

[tamasgal]

Ahm, yes, the numba_jim version which is not using the builder is the clear winner for an array size of arround 8000 and above :) I also checked if the results are the same and they are.

[jpivarski]

The reason the numba_jim version is flat with respect to "size of array" is because it is computing it for all elements every time. (That is, the "size of array" is always maximal; slicing arr[0:n] doesn't remove data from the layout.)

I've fixed the implementation so that it ignores inaccessible data on the high side (scales appropriately for arr[:n] but not arr[-n:]). The changed lines are labeled "FIX". The main benefit is that it makes the plot right.

@nb.njit
def base_case(arr, start, stop):
    out = np.empty(len(arr), np.bool_)
    for i, el in enumerate(arr):
        out[i] = len(el) > 0 and el[0] == start and el[-1] == stop
    return out

def mask_startstop(arr, start, stop):
    def recurse(layout):
        if layout.purelist_depth == 2:
            np_array = base_case(ak.Array(layout), start, stop)
            return ak.layout.NumpyArray(np_array)

        elif isinstance(layout, (
            ak.layout.ListArray32,
            ak.layout.ListArrayU32,
            ak.layout.ListArray64,
        )):
            if len(layout.stops) == 0:                                   # FIX
                content = recurse(layout.content)                         # FIX
            else:                                                         # FIX
                content = recurse(layout.content[:np.max(layout.stops)])  # FIX
            return type(layout)(layout.starts, layout.stops, content)

        elif isinstance(layout, (
            ak.layout.ListOffsetArray32,
            ak.layout.ListOffsetArrayU32,
            ak.layout.ListOffsetArray64,
        )):
            content = recurse(layout.content[:layout.offsets[-1]])        # FIX
            return type(layout)(layout.offsets, content)

        elif isinstance(layout, ak.layout.RegularArray):
            content = recurse(layout.content)
            return ak.layout.RegularArray(content, layout.size)

        else:
            raise NotImplementedError(repr(arr))

    layout = ak.to_layout(arr, allow_record=True, allow_other=False)
    return ak.Array(recurse(layout))

[tamasgal]

OK, I thought it was some other overhead. I’ll rerun it tomorrow! Thanks already, I learned a lot.

[tamasgal]

I confirm, it's linear now and the fastest:

[tamasgal]

I think we can close this, thanks Jim for the fruitful discussion and lots of input (including the final implementation prototype)!

EDIT: ah, I just realised it's now a Q&A ;)

[jpivarski]

I've migrated the issue into a Discussion, which doesn't need to be closed. It gets to remain available to help other users without confusing me about what work needs to be done. That's why I used to recommend StackOverflow, but since GitHub Discussions are all under one roof, I can convert Issues into Discussions, like this one.

[tamasgal]

Yes! I guess I should have started the #584 also in stackoverflow. I can move it if you want to!

[tamasgal]

Btw. what I do not like in StackOverflow is that there is no easy back-and-forth communication. Maybe we can simply copy over a condensed version of the discussion here to SO after an answer has been figured out ;)

[jpivarski]

If GitHub Discussions work for us, I'll stop recommending StackOverflow. I'm already promoting them on an equal level (with GitHub Discussions first). The ability to turn Issues into Discussions is a big plus—the problem had been that people put the wrong types of questions on Issues and StackOverflow, and then it was stuck there. ("Wrong" meaning "that which should stay open after being answered, to help other users" versus "that which should be closed after the issue has been fixed.") On the other hand, consolidating everything into the hands of one company because it's easier is a slippery slope toward monopoly (remembering that GitHub is owned by Microsoft).

Beyond just the advantage of being under one roof, GitHub Discussions approached the problem of back-and-forth differently: there are no upvotes to change the order, so it stays mostly linear, and it has this two-level hierarchy of nesting instead of StackOverflow's tiny comment boxes (tiny because they're discouraged). On the one hand, that might make it easier to find the right answer in StackOverflow because the right stuff is sifted up and the wrong or less useful stuff is sifted down, but if there's any kind of back-and-forth to be had, GitHub Discussions is better. In Awkward Array and Uproot questions, back-and-forth has generally been useful. (There aren't so many "how do I append to a list?" questions.)