Flattening nested RecordArrays and restoring records from flattened Array

[kreisl]

Hi,

I am analyzing a nested record structure shaped similar to:

array = ak.Array([
    {
    "a": {"b": [1,1], "c":[2,2]},
    "d": {"b": [3,3], "c":[4,4]}
    },
    {
    "a": {"b": [5,5], "c":[6,6]},
    "d": {"b": [7,7], "c":[8,8]}
    }
])

What is the recommended way of unraveling the array such that all the record are flattened like so:
ak.Array([[1,1,2,2,3,3,4,4], [5,5,6,6,7,7,8,8]]) or ak.Array([[1,1,3,3,2,2,4,4], [5,5,7,7,6,6,8,8]]) which I need for processing inside pytorch.
Right now I do the following:

def unzip_and_concatenate_test(array, labels):
    array = ak.concatenate(ak.unzip(array[labels.pop(0)]), axis=0)
    if labels:
        return unzip_and_concatenate(array, labels)
    return array

where the labels are given by for example [['a','d'],['b','c']].
Is there a simpler approach to flattening nested records?

I also have another related question. After I flattened the array is it possible to "view" it as records again, or would I need to copy and create a new Array using the Arraybuilder? I may have to mask certain entries (for example masking all entries from "c") of the flattened array after processing, which would be easier when the array would be viewable as a record rather than computing the positions by hand.
Below I have sketched the workflow:
RecordArray -> flattened Array -> processing in pytorch -> masking the Array -> postprocessing

Thank you and best regards,
Lukas

[agoose77]

On the face of it, doing things by hand seems like the easiest approach. We do have tools to generalise this for arrays, but it involves constructing layouts yourself. Do you have any more information about the structure, e.g. are your arrays regular?

[kreisl]

My arrays are regular. Every entry in the array has the same shape: for example

[{
  "class_1": {"property_1": [1,2], "property_2":[1,2,3]},
  "class_2": {"property_1": [1,2], "property_2":[1,2,3]}
},{
  "class_1": {"property_1": [4,3], "property_2":[1,2,6]},
  "class_2": {"property_1": [0,1], "property_2":[1,4,3]}
}]

flattened array:

[ [1,2,1,2,3,1,2,1,2,3], [4,3,1,2,6,0,1,1,4,3] ]

Would it be possible to view this flattened array again with the "records" layout? This would make post processing easier (for example selecting only the property_2 of class_1).

[agoose77]

For some context, Awkward is not really designed to be able to dissolve fields - they're an important part of the structure. For this reason, we don't have any built-in routines to do this.

If you have a small number of fields (e.g. 2 properties, 2 classes), then manually concatenating is fine. To rebuild the structure, you can use ak.zip with strided views over the ravelled array, e.g.

res = ...

b = res[:, :2]
c = res[:, 2:4]
b2 = res[:, 4:6]
c2 = res[:, 6:8]

array_res = ak.zip(
    {"a": ak.zip({"b": b, "c": c}, depth_limit=1), "d": ak.zip({"b": b2, "c": c2}, depth_limit=1)}
)

[jpivarski]

The arrays are always fixed-length and you want to go between views of them as a single contiguous array and as records with fields, that's sounding more like NumPy structured arrays than Awkward record arrays.

In a NumPy structured array, the memory layout is a single buffer in which fields must be contiguous, equivalent to an array of

struct {
    struct {
        int b[2];
        int c[2];
    } a;
    struct {
        int b[3];
        int c[3];
    } d;
};

in C (I don't remember the exact syntax, but I think you understand what I mean). This data can be viewed ("reinterpret cast") as a single array of integers, or as an array of one-level structs. Views are fast, just interpreting the in-memory buffer in different ways.

In an Awkward record array, the memory layout is separate for each field. array.a.b may be in a completely different memory address from array.a.c, and each of them are one-dimensional arrays of integers. That makes other types of rearrangement fast: incorporating array.a.b in some other_array, splitting all fields out (ak.unzip), combining fields into new records (ak.zip), etc. Those are just interpreting the in-memory buffers (plural) in different ways.

However, turning an Awkward record array into a single flat array without fields is not a trivial operation. You'd do it by picking fields out and calling ak.concatenate on them, possibly with an explicit axis. Pulling the pieces back into fields would involve the slices @agoose77 specified. Slices are just views (fast), but you have to be very aware of the regular lengths to slice them correctly.

Given that you want to do the records-to-single-array and back frequently, NumPy's data structure sounds more appropriate to your problem. It is possible to make nested fields in a NumPy structured array; I've recently learned about that, though I haven't tried it.

Let's see...

>>> import numpy as np

>>> flat = np.array([1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8])

>>> structured = flat.view([("a", [("b", int, 2), ("c", int, 2)]), ("d", [("b", int, 2), ("c", int, 2)])])

>>> structured
array([(([1, 1], [2, 2]), ([3, 3], [4, 4])),
       (([5, 5], [6, 6]), ([7, 7], [8, 8]))],
      dtype=[('a', [('b', '<i8', (2,)), ('c', '<i8', (2,))]), ('d', [('b', '<i8', (2,)), ('c', '<i8', (2,))])])

>>> structured["d"]["b"]
array([[3, 3],
       [7, 7]])

>>> q = structured["d"]["b"]
>>> q[0, 0] = 999    # mutable and still linked to the original
>>> q
array([[999,   3],
       [  7,   7]])
>>> structured       # yup, still linked
array([(([1, 1], [2, 2]), ([999,   3], [4, 4])),
       (([5, 5], [6, 6]), ([  7,   7], [8, 8]))],
      dtype=[('a', [('b', '<i8', (2,)), ('c', '<i8', (2,))]), ('d', [('b', '<i8', (2,)), ('c', '<i8', (2,))])])

That works, and this structured array is re-viewable into a flat one and back.

Having this single buffer is an advantage and a disadvantage. If you really want the fields to be "dissolvable" into a single array (of a specific, particular order), then it's just what you want. Awkward records weren't designed like this because the idea is that record fields describe semantically distinct data, like columns with different names, data types, and probably different units of measurement. Those are values that you'd never want to mix indiscriminately, but you likely want to pull a column out and deal with it separately from the rest of the structure.