promote_option.py

def promote_option(
    array: ak.Array,
    axis: int = None,
    union: bool = True,
    highlevel=True,
    behavior=None,
) -> ak.Array:
    nplike = ak.nplike.of(array)
    layout = ak.to_layout(array)
    posaxis = layout.axis_wrap_if_negative(axis)

    def getfunction(layout, depth):
        if depth != posaxis + 1:
            return

        if not isinstance(layout, ak.layout.RecordArray):
            return

        # Compute union of mask
        invalid = nplike.asarray(layout.contents[0].bytemask())
        if union:
            for content in layout.contents[1:]:
                nplike.bitwise_or(
                    invalid, nplike.asarray(content.bytemask()), out=invalid
                )
        # Otherwise, require all masks to be identical
        else:
            for content in layout.contents[1:]:
                if not nplike.all(nplike.equal(invalid, content.bytemask())):
                    raise ValueError
        invalid = invalid.astype(np.bool_)

        # Project mask from all contents
        mask = ak.layout.Index8(invalid)
        new_contents = [c.project(mask) for c in layout.contents]

        # Build new RecordArray with only projected contents
        new_layout = ak.layout.RecordArray(
            new_contents,
            layout.recordlookup,
            len(new_contents[0]),
            layout.identities,
            layout.parameters,
        )

        index = nplike.cumsum(~invalid, dtype=np.int64) - 1
        index[invalid] = -1
        out = ak.layout.IndexedOptionArray64(ak.layout.Index64(index), new_layout)
        return lambda: out

    out = ak._util.recursively_apply(layout, getfunction)
    return ak._util.maybe_wrap_like(out, array, behavior, highlevel)