getting around the literal string requirement for indexing

[ananis25]

Hi - awkward arrays have made processing arrow datasets with numba really nice! However, I can't quite figure how to use the combo for more dynamic use cases.

I have a dataset of the type - N * {"a": int64, "b": int64, "events": var * {"c": int64, "d": int64}}. Essentially, a list of sequences where each sequence contains some attributes and a list of events, which themselves have event level attributes.
I am trying to write a regex style pattern matcher, where we match against a stream of events rather than characters. The numba routine goes over each row of the dataset and checks if the events array has a consecutive set satisfying the pattern condition.

Pattern matching is recursive, so it is hard to statically determine when a specific event attribute is accessed. And I get a lot of errors of the kind,TypeError("only a *constant* field name string may be used as a record slice in compiled code").Issue #1420 makes it clear that non literal strings can't be used to index an awkward array since the output could be a scalar/array of any type.

However, given the events attribute is a homogenous record array, with only integer values, is it possible to override how awkward types the getitem call for it? All attribute access code is of the form - dataset[i]["events"][j][prop], where prop is a variable string.

[agoose77]

How many fields (columns) do you have? One easy solution would be to explode the records into a list of arrays, which can then be indexed, e.g.

event_fields = ak.unzip(events)

def process_events(event_fields):
    j = choose_column_index()
    do_something_with(
        event_fields[j]
    )

As @jpivarski notes in #1420, supporting non-literal strings is an intended feature, but it's not the highest priority right now with all of the other things that are being worked on!

[ananis25]

Got it, tuples can still be heterogenous so integer indexing can't guarantee that output type will always be the same. I took your suggestion to construct a 2D array instead! Hopefully, this is sufficient for what I need.

import awkward as ak

data = {
    "a": np.arange(100),
    "b": np.arange(100, 200),
    "events": [
        [{"c": j if j == 0 else None, "d": j + 1} for j in range(i)]
        for i in range(200, 300)
    ],
}
awk_data = ak.Array(data)

events = awk_data["events"]
list_arrays = [events[f] for f in events.fields]
list_counts = ak.num(list_arrays[0])

# to stack the arrays for event-level attributes, we need them to be 2d
list_arrays_2d = [ak.unflatten(ak.flatten(arr), 1) for arr in list_arrays]
list_array_flat = ak.to_regular(ak.concatenate(list_arrays_2d, axis=1))
out_array = ak.unflatten(list_array_flat, list_counts)

print(events.type)
# 100 * var * {"c": ?int64, "d": int64}

print(out_array.type)
# 100 * var * 2 * ?int64

generate a function that is then jitted. You can do this with AST generation or string exec, whether you prefer!

I was indeed dreading having to create functions by composing strings! Could you please link to any examples of AST generation for numba code? Thank you.

[agoose77]

The reason that I don't think (Awkward) tuples will work is that they are still implemented using string indices, so the same restrictions apply as with records. Furthermore, the .slot<N> attributes don't work in Numba (IIRC), so they're no use here either.

You might find it slightly cleaner to write list_arrays = ak.unzip(events).

If you pass in list_arrays directly to your Numba function, it is likely to be more performant than stacking them into a regular 2D array, because if they're contiguous (which they sometimes may be if you've not transformed them) the Numba conversion routines will not need to copy them. That said, I can't remember whether we check this or just make a copy regardless, though I doubt it.

You can pass in a separate 2D empty array (or construct it in the function), because you'll need to allocate memory for the result.

Here's some pseudo code that illustrates what I mean:

list_arrays = ak.unzip(events)
result = np.empty((len(events), len(events.fields)), dtype=np.float64)

@njit
def transform(list_arrays, result):
    first = list_arrays[0]
    
    for i in range(len(first)):
        j = select_some_column(list_arrays, i)
        result[i, j] = ...
        
transform(list_arrays, result)

[ananis25]

If you pass in list_arrays directly to your Numba function, it is likely to be more performant than stacking them into a regular 2D array, because if they're contiguous (which they sometimes may be if you've not transformed them) the Numba conversion routines will not need to copy them. That said, I can't remember whether we check this or just make a copy regardless, though I doubt it.

Yeah, avoiding a copy would have been better, though I need to multiindex to keep the logic simple. So, a single array of type N * var * M * int64 is preferable to a list of M arrays of the type N * var * int64. Numba gives me grief with reflected lists too.
I will explore writing a wrapper class akin to awkward for the list of arrays to get around copying.

[agoose77]

@ananis25 here's an example of AST generation:

https://gist.github.com/agoose77/ec37316c72214b47ebf06b1956a72072

I'm fairly confident that there is a way to do this using Numba's API, but I'm not yet familiar enough with Numba's internals to be able to suggest a solution in the time I have available to look at this problem :) AST generation is an "obvious" way to solve this, even if it's a bit clunky.

[ananis25]

No worries, thank you for taking the time! I guessed you were hinting at something like numba rewrite passes, though using the python ast module also gave me an idea about how to write filter functions for events.