concatenate arrays of records without getting Unions

[ivirshup]

I would like to concatenate arrays of records without getting unions back. Here's a quick illustrative example of a use case:

import awkward._v2 as ak

awk_a = ak.Array([
    {"a": [1, 2, 3], "b": [1, 2]},
    {"a": [4, 5], "b": [3, 4]},
    {"a": [6], "b": [5]},
])


awk_b = ak.Array([
    {"a": [1, 2, 3]},
    {"a": [4, 5]},
    {"a": [6]},
])

ak.concatenate([awk_a, awk_b])

I would like to end up with a non-Union result here. E.g. something like:

# '6 * {a: var * int64, b: var * int64}'
ak.Array([
    {"a": [1, 2, 3], "b": [1, 2]},
    {"a": [4, 5], "b": [3, 4]},
    {"a": [6], "b": [5]},
    {"a": [1, 2, 3], "b": []},
    {"a": [4, 5], "b": []},
    {"a": [6], "b": []},
])

# '6 * {a: var * int64, b: option[var * int64]}'
ak.Array([
    {"a": [1, 2, 3], "b": [1, 2]},
    {"a": [4, 5], "b": [3, 4]},
    {"a": [6], "b": [5]},
    {"a": [1, 2, 3]},
    {"a": [4, 5]},
    {"a": [6]},
])

I would ideally like to be able to do this at concatenation time (and even be able to choose just to concatenate the intersection of keys). But, being able to convert from a union to an optional type would also work.

I'm having some trouble finding docs for this, how would I go about doing this?

[agoose77]

I'd suggest transforming awk_b such that it has the necessary field. You could make use of the ak.Array.__setitem__ feature here, which allows you to add fields to an existing array (advanced users note that you can modify arrays, but not layouts, so this feature replaces the existing layout with a new one). The nice benefit of using __setitem__ (or indeed ak.with_field) is that it broadcasts, so by doing

awk_b['b'] = None

you end up with

>>> awk_b.type.show()
3 * {
    a: var * int64,
    b: ?unknown
}

Upon concatenation, this then becomes:

>>> ak.concatenate([awk_a, awk_b]).type.show()
6 * {
    a: var * int64,
    b: option[var * int64]
}

Awkward is still creating a union for the different layout nodes under the hood; unknown translates to EmptyArray in Awkward's layout node structure, whilst var * int64 is a list-type over a NumPy array. Yet, Awkward has rules about which types can be merged, and which require a union, and EmptyArray can always be merged.

[ivirshup]

From the other conversation, do you think ak.transform could be a good way to implement this?

That way I could recurse down a set of arrays and "broadcast" them to an inner or outer union before concatenation.

[jpivarski]

The __setitem__ that @agoose77 suggested,

At present the easiest ways to do this are to take an existing array and e.g. slice it to zero length

awk_b['b'] = awk_b['a'][..., :0]

does use broadcasting to create the new field. (It's implemented with the same backend as ak.transform.) Since this high-level mechanism exists to do it, this is definitely simpler than rolling your own.

[ivirshup]

Just trying this out now, and running into an issue. I can get the expected result from:

awk_a = ak.Array([
    {"a": [1, 2, 3], "b": [1, 2]},
    {"a": [4, 5], "b": [3, 4]},
    {"a": [6], "b": [5]},
])


awk_b = ak.Array([
    {"a": [1, 2, 3]},
    {"a": [4, 5]},
    {"a": [6]},
])

ak.concatenate([awk_a, ak.with_field(awk_b, None, "b")]).type.show()

6 * {
    a: var * int64,
    b: option[var * int64]
}

However, I get a union from:

awk_c = ak.Array([
    [{"a": 1, "b": "foo"}],
    [{"a": 2, "b": "bar"}, {"a": 3, "b": "baz"}]
])

awk_d = ak.Array([
    [{"a": 4}, {"a": 5}],
    [{"a": 6}],
    [{"a": 7}],
])

ak.concatenate([awk_c, ak.with_field(awk_d, None, "b")]).type.show()

5 * union[
    var * {
        a: int64,
        b: string
    },
    var * {
        a: int64,
        b: ?unknown
    }
]

Is this expected?

[ivirshup]

As an aside, is there a way to do something like:

awk_c.astype(ak.types.from_datashape("2 * var * {a: int64, b: ?string}"))

If not, is there a fundamental barrier to something like this being possible?

[jpivarski]

Is this expected?

I think you found a bug: we want "string" and "option[unknown]" to unify to "option[string]", and it's not recognizing that. It should first notice that one of the two is "option" and make the result type "option", then unify "string" (or "X") with "unknown" to get "string" ("X"), since "unknown" is an identity in type-unification. I think this is supposed to happen somewhere in simplify_uniontype. It's interesting that it got it right when "X" is a primitive type, int64, but not when it's a list type. This is something we should promote into an issue.

As for casting an entire array as an entire type, that sounds like a good function to have, though it would take some thinking to do it right. You'd have to recurse down the array tree and the type tree, modifying array elements based on what we see in the type, and it wouldn't be one-to-one because high-level types are simpler than arrays. So, it's not a fundamental barrier, but it will be quicker for us to fix the bug in type-unification, above.