Getting size of a given dimension

[grst]

I would like to follow up on a discussion on gitter a while ago. Since then I have understood a bit better what I actually need and, while I'm still only scratching the surface, understood a bit better how awkward arrays work.

What I need is a function like this:

def dim_len(array, dim):
   """
   Get the length of the array in the given dimension.

   Returns a non-negative integer if the array is regular in the given dimension and None otherwise. 

   Examples
   --------
   >>>  dim_len(ak.Array(np.array([[1,2,3],[3,4,5]])), 1) 
   3
   >>> dim_len(ak.Array([[1,2,3],[3,4,5]])), 1)
   None
   >>> dim_len(ak.Array([{'a': 1}, {'b': 1}]), 0)
   2
   """
   raise NotImplementedError

In the chat on gitter, it was suggested to recurse through the array types. Is this the way to go? It feels a bit weird that no such function is available as part of the library, I would have expected this to be a rather common use-case.

If using the array types is the way to go, could I please get your feedback on if I understood correctly the implications of each type:

• ArrayType: outermost type, retrieve .length attribute
• ListType: variable-length
• NumpyType: regular type, retrieve .size attribute
• OptionType: Does not represent an own dimension. Get the .size attribut of the child in .content.
• RecordType: End of the recursion. Dimensions after a record type are not counted.
• RegularType: regular type, retrieve .size attribute
• UnionType: Does not represent an own dimension. Get .size attribute of the children in .contents?. Could it be that different children have different sizes? Can there be a regular dimension after a UnionType?
• UnknownType: ??

CC @ivirshup

[agoose77]

Your table is mostly correct. Whilst I was thinking about this, I noticed that I initially had some confusion as to how types compare with layouts.

The notable difference between types and layouts is that in the type system, UnknownType and NumpyType represent atoms that have no length. Whereas, with layouts, the smallest unit is a singular array (NumpyArray or EmptyArray). The dimension to which the types belong is their parent, e.g. RegularType, ListType, etc., whereas the layouts list-types are the dimension. Hence, if you encounter UnknownType or NumpyType (atoms with no length) during recursion you've gone too far!

Meanwhile, the list-types, e.g. RegularType, ListType, correspond to inner-dimensions; the root ArrayType is always present, and represents the outermost dimension.

I wrote out an example of dim_len, exploding it into different routines to make it clear what is happening at each level:

import awkward._v2 as ak
import typing as t

SizeType = t.Union[int, None]


def _dim_size_next(node: ak.types.Type, depth: int) -> SizeType:
    """Walk into the contents of the given node.

    :param node: type node
    :param depth: depth above the intended type node
    :returns: size of type, or None if irregular
    """
    # Do we have a container?
    if isinstance(node, (ak.types.ListType, ak.types.RegularType, ak.types.ArrayType)):
        return _dim_size_impl(node.content, depth - 1)
    else:
        raise TypeError(f"Unexpected node type: {node!r}")


def _dim_size_at(node: ak.types.Type) -> SizeType:
    """Compute the size of the given type node

    :param node: type node
    :returns: size of type, or None if irregular
    """
    if isinstance(node, ak.types.RegularType):
        return node.size
    elif isinstance(node, ak.types.ListType):
        return None
    elif isinstance(node, ak.types.ArrayType):
        return node.length
    else:
        raise TypeError(f"Unexpected node type: {node!r}")


def _dim_size_transform(node: ak.types.Type) -> ak.types.Type:
    """Transform the given type node to ignore superfluous types
    
    :param node: type node
    :returns: important type
    """
    # These types are "atoms", and do not correspond to a dimension.
    # As atoms, they have no length, so if we see them, it's an error.
    if isinstance(node, (ak.types.NumpyType, ak.types.UnknownType)):
        raise ValueError(f"Cannot compute size of atom: {node!r}")
    # Meanwhile, unions *could* have an interpetable size, but we choose 
    # not to handle them. Records are never considered to have a size
    elif isinstance(node, (ak.types.UnionType, ak.types.RecordType)):
        raise TypeError(f"Cannot compute size through branching types: {node!r}")

    # Whether recursing through, or looking at this dimension
    # We don't care about options, so move through them
    if isinstance(node, ak.types.OptionType):
        return _dim_size_transform(node.content)
    else:
        return node


def _dim_size_impl(node: ak.types.Type, depth: int) -> SizeType:
    """Dispatcher to compute the size of the given type node

    :param node: type node
    :param depth: depth above the intended type node
    :returns: size of type, or None if irregular
    """
    # Depth-agnostic transforms
    node = _dim_size_transform(node)

    # Now consider whether we're recursing or at the appropriate depth
    if depth > 0:
        return _dim_size_next(node, depth)
    else:
        return _dim_size_at(node)


def dim_size(array: ak.Array, dim: int) -> SizeType:
    """Compute the size of a particular dimension of the given array

    :param array: Awkward Array
    :param dim: dimension of which to compute the size
    :returns: size of dimension, or None if irregular
    """
    return _dim_size_impl(array.type, depth=dim)

Awkward Array operations generally stop at records, so it doesn't make sense to compute the size of these types. However, it is technically possible to compute the size of unions; you just need to resolve the branches at the end. I opted to just make this an error, but you could extend it.

Note that this doesn't handle negative axis (dim) values. You could do this by first computing the depth of the array, e.g. via ndim. Again, this needs special handling for unions; ndim ignores the union dimension in these cases.

You might use this like:

import numpy as np
arr = ak.Array([
    [
        [
            [1, 2, 3],
            [4, 5, 6, 7]
        ],
        [
            [1, 2, 3],
            [4, 5, 6, 7]
        ]
    ]
])
arr = ak.to_regular(arr, axis=2)

print(dim_size(arr, 0))
print(dim_size(arr, 1))
print(dim_size(arr, 2))
print(dim_size(arr, 3))

[jpivarski]

As usual, @agoose77 beat me to an answer again! (Sarcasm: I don't mind!)

I was writing an answer that uses the new ak.transform (PR #1610), which I think was inspired, or at least got bumped up in priority by your problem. (It fulfills an old issue #516.) The goal of this interface was to provide a public API for the function we use internally to define functions like these, which would streamline the process of absorbing it into the Awkward codebase, once we have a good idea of what the general API should look like.

Here's what I came up with, and most of the complication is dealing with the fact that union types introduce branches.

def dim_len(array, axis):
    if axis < 0:   # negative axis is another can of worms... maybe later
        raise NotImplementedError
    elif axis == 1:
        return len(array)
    else:
        def size_at_depth(layout, depth, lateral_context, **kwargs):
            if layout.is_NumpyType:
                # if it's an embedded rectilinear array, we have to deal with its shape
                # which might not be 1-dimensional
                if layout.is_UnknownType:
                    shape = (0,)
                else:
                    shape = layout.shape
                numpy_axis = lateral_context["axis"] - depth + 1
                if not (1 <= numpy_axis < len(shape)):
                    raise TypeError(f"axis={lateral_context['axis']} is too deep")
                lateral_context["out"] = shape[numpy_axis]
                return layout.nplike.empty(1)

            elif layout.is_ListType and depth == lateral_context["axis"]:
                if layout.is_RegularType:
                    # if it's a regular list, you want the size
                    lateral_context["out"] = layout.size
                else:
                    # if it's an irregular list, you want a null token
                    lateral_context["out"] = -1
                return layout.nplike.empty(1)

            elif layout.is_RecordType:
                # if it's a record, you want to stop descent with an error
                raise TypeError(f"axis={lateral_context['axis']} is too deep, reaches record")

            elif layout.is_UnionType:
                # if it's a union, you could get the result of each union branch
                # separately and see if they're all the same; if not, it's an error
                result = None
                for content in layout.contents:
                    context = {"axis": lateral_context["axis"]}
                    ak.transform(size_at_depth, content, lateral_context=context, return_array=False)
                    if result is None:
                        result = context["out"]
                    elif result != context["out"]:
                        raise TypeError(f"union results in different values at axis={lateral_context['axis']}")
                lateral_context["out"] = result
                return layout.nplike.empty(1)

        # communicate with the recursive function using a context (lateral)
        context = {"axis": axis}

        # "transform" but we don't care what kind of array it returns
        ak.transform(size_at_depth, array, lateral_context=context, return_array=False)

        # you wanted the null token to be None
        return None if context["out"] == -1 else context["out"]

The primary purpose of ak.transform is to turn one array into another array, leaving all the node types that we don't care about intact. In your case, however, we don't care about the output array, so we pass return_array=False and terminate recursion by returning a dummy array, np.empty(1) (using layout.nplike instead of np so that it works on all backends, CuPy, etc.).

We care about NumpyType arrays that are the leaves of a tree, as they might have their own rectilinear shape that may be part of the answer. The UnknownType is for leaves whose dtype is unknown because they came from an untyped source (e.g. JSON) with no examples to disambiguate the type. So their shape is (0,) by definition. We could have left out the whole layout.is_NumpyType section by passing numpy_to_regular=True in ak.transform, so that the leaves would only ever be 1-D and all that rectilinear shape information would be in RegularArrays, rather than NumpyArray, but that would unnecessarily force non-contiguous arrays to have to be rewritten. All you want is the dimension, after all.

We care about ListType arrays because that's where the answer will usually come from. If it's regular, you get the size, and if it's not, you want a null token.

We care about RecordType because you want the recursion to stop there (which it wouldn't normally do).

We care about UnionType because you either want to stop there, as in @agoose77's solution, or you want to unify the results of each branch, which is possible. This is why I made the null token temporarily -1, so that it could be distinguished from None.

Stepping back, it's a fairly complex function (a page), but it works for every possible Awkward Array.

Edit: Oh, I forgot to mention that we didn't have to mention option-types at any point, because they're irrelevant for this problem. ak.transform lets you skip past them. (When size_at_depth returns None, the recursion keeps going.)

[ivirshup]

Alright, thanks for the extra details

we don't have an ak._v2.types.IndexedType, for example

Is this the reason for having all these attributes instead of going through instance checks? It seems like it has effectively been made user visible through the attributes.

[agoose77]

The array.type attribute generates an array type from the array form, which is our serialisable layout metadata that combines with buffers to build a layout. So, at the layout level, although we can generate these type objects, they're generated lazily on demand from the form.

The is_XXXType flags are really just a quality of life improvement for reading (and writing) source code, as well as a slight performance boost (I suspect).

My advice in terms of thinking about these things is to treat the array.type attribute as a convient way to interact with the abstract user-facing structure of the array. For anything more detailed / advanced, go straight to the layouts (which are a richer source of information).

[jpivarski]

What I didn't recognize when I started writing a solution using ak.transform is that this particular problem can be solved by looking at types alone, as @agoose77 did. And that solution is much simpler; I hadn't thought of it. The only possible down-side of having two ways of doing something, one of which can only be applied in a subset of cases, is if you later want to change the function beyond that subset, you wouldn't be able to do so by a small extension, you'd have to change your whole technique.

Suppose, for instance, that instead of returning None for variable-length lists, you later want to return the minimum length or the maximum length. That information isn't in the type; you'd have to switch from recursing over types to recursing over arrays, and the way of doing that is entirely different, so it would mean starting over.

If you're certain that you'll only need type information, I think @agoose77 is right: the type-only recursion is much simpler. That's what I'd go for, if I'd thought of it first.

Why not plain recursion?

You could recurse over the array layout nodes (subclasses of ak._v2.contents.Content) in the same way that you recurse over types (subclasses of ak._v2.types.Type). If that recursive function is going to be applied to arbitrary Awkward Arrays, it would just need to know the full set of Content subclasses, just as it now needs to know the full set of Type subclasses (neither of which is expected to change in the future). This function doesn't care about option-types: it wants to just recurse through them if it finds any in the tree, so it needs to catch the

isinstance(type_node, OptionType)

so that it can continue down type_node.content, or similarly

isinstance(layout_node, (BitMaskedArray, ByteMaskedArray, IndexedOptionArray, UnmaskedArray))

so that it can continue down layout_node.content. You would need to know the full set of types to check for and whether it has zero children (EmptyArray and NumpyArray), multiple children (RecordArray and UnionArray), or one child (everything else). Fortunately, the names of child/children attributes are systematic: content or contents for all layout nodes, type nodes, and Form nodes (explanation below).

So, you could do this, and we started with this in the early days of writing functions on Awkward Arrays. In fact, the earliest functions were implemented as a method on all Content subclasses, following proper OOP, but this spread the definition of a single function across many files and we will not be adding new Content subclasses. So then we were writing these functions that did isinstance for every Content subclass, but most subclasses were just a pass-through and there was a lot of redundant boilerplate. Furthermore, functions of more than one array had to consider cases in which the array types didn't match at a given level, e.g. one is a ListArray and the other is a ListOffsetArray, and this mixes broadcasting with the calculation of the operation we're actually interested in, so we clearly had to separate the process of recursing over the tree (with possible broadcasting) from the calculation of the operation we're actually interested in.

This visitor pattern is what we came up with; actually, it's a rewrite of the original visitor pattern with a cleaned-up interface. And then we made it public through ak.transform so that other libraries can get the same benefits. So while you could write a direct recursive function over the layout nodes, you'd find that ak.transform would reduce boilerplate. The example above doesn't mention option-types, for instance.

Why the is_*Type attributes?

The next part is the is_*Type attributes. Yes, they're just to reduce typing and screen clutter:

isinstance(layout_node, (BitMaskedArray, ByteMaskedArray, IndexedOptionArray, UnmaskedArray))

becomes

layout_node.is_OptionType

Maybe it's marginally faster to execute, but that wasn't the motivation.

Relationship between Content types, Type types, Form types (for completeness) and is_*Type attributes

Inside of an ak.Array's layout, a tree of Content subclasses constitute the array itself. It would be possible to answer any question by examining this tree.

But several different Content subclasses correspond to the same high-level Type, and these are presented to the user as though they were the same. For instance, ListArray and ListOffsetArray are two ways of representing variable-length lists, the first is more versatile, the second is amenable to more optimizations, but a high-level user only cares that they have an array of lists. So when you invoke ak.Array.type, it creates a Type tree from the Content tree by simplifying/projecting away information. This Type is usually viewed as a Datashape string, which includes the length of the array, so the top-most element in this hierarchy is an ArrayType containing the length. (And issubclass(ArrayType, Type) is False: it's outside the hierarchy, similar to the way that ak.Array is not a Content subclass.)

For completeness, I should mention that we sometimes need low-level types, which are called Forms for a word like "type" that isn't "type" (or "sort", and "kind" has a different meaning in type theory). These are one-to-one with Content subclasses; they only lack the array data itself. It's used, for instance, as a template for constructing arrays in ak.from_buffers).

The is_*Type attributes were named "Type" because they mostly correspond to high-level Types. But their exact definitions came from what combinations of Content subclasses we frequently found ourselves checking in isinstance, which has some overlaps, making it not a strict hierarchy, the way that actual Types are. For instance,

layout node	is-attributes	Form	Type
ListArray	is_ListType	ListForm	ListType
ListOffsetArray	is_ListType	ListOffsetForm	ListType
RegularArray	is_ListType, is_RegularType	RegularForm	RegularType

because a lot of the calculations we wanted to apply were the same for all three list-like types, but some applied differently to regular lists.

Also, some Content elements correspond to no type: an IndexedArray containing a Content of type T has type T. When projecting the information out to make a Type tree, these nodes disappear entirely. However, there's an is_IndexedType attribute to detect it when recursing over layout nodes (as @agoose77 pointed out).

So maybe the "Type" part of the is_*Type attribute names should be changed. We still have time to do so: the is_*Type attributes are a v2 addition and v2 hasn't been finally/formally released yet. But if not "Type," then what? Since they overlap, they look to me like type classes in Haskell, but "Class" would be a pretty bad name, given that this is in Python. Maybe is_*TypeClass?

Why are there so many ways of describing types in Awkward Array?

I haven't even mentioned the hierarchies of type objects that enable iteration in Numba or JIT-compilation in C++ Cling. Considering these 2 and the above 3, that's 5 different systems for describing the type of an Awkward Array. That's why we're not going to be adding any new node types—there's so much we'd have to change. It's a closed system, which is why it was important for it to be a complete system to start with (i.e. not just ragged arrays of numbers).

[grst]

I'm a bit confused by the string representation, which results in two nested types:

>>> ak.Array(["foo", "foo", "bar"]).type
ArrayType(ListType(NumpyType('uint8', parameters={'__array__': 'char'}, typestr='char'), parameters={'__array__': 'string'}, typestr='string'), 3)

compared to

>>> ak.Array([1, 2, 3]).type
ArrayType(NumpyType('int64'), 3)

This is also an edge case for the dim_len implementation above:

# expected: raise TypeError
# actual: 
>>> dim_len(ak.Array(["foo", "foo", "bar"]), 1) 
None

Why can't the string array be a single NumpyType like its int counterpart? And where would you suggest it's best to handle this edge-case in the dim-len function?

[jpivarski]

Regardless of how it's represented, an array of strings has to be implemented like an array of lists of uint8; the variable-length lists is the part that requires special care. (It can't be a NumpyType because NumPy doesn't even attempt to deal with variable-length strings.) In an earlier version of Awkward Array (version 0.x), strings were implemented as a special node type, separate from JaggedArray, leading to a lot of duplication of code.

Our point of view on strings is that they're physically like lists of uint8, in how they're laid out in memory (and disk) and how they're processed in slicing and other basic operations, but they have special high-level behaviors: their repr shows them as letters in quotation marks, == compares whole strings as atomic units, not their character contents the way a list would, and so on. It's as though our strings are a subclass of lists, a specialization with overridden methods. Overriding records in Awkward Array is very common in HEP, so we're using the same ak.behavior mechanism to do both. (At least ideally; strings have some hard-coded specializations because they're so special—we're figuring out what to do about that.)

To handle it in dim_len, you can take a Content, Form, or Type subclass and ask

node.parameter("__array__") in ("string", "bytestring")

if you want to exclude UTF-8 strings and unencoded bytestrings in particular, or

node.parameter("__array__") is not None

if you want to exclude any overridden array types at all. If a parameter is not defined on a node, this function will return None (like a dict's get), so undefined and defined-to-be-None can be treated the same way.