Protocol extensions for awkward arrays

[alimanfoo]

This issue is a starting point for discussing possible protocol extensions to allow support for various types of "awkward" arrays that do not fit into the model of arrays with fixed size dimensions and/or with fixed size data types.

For example, this includes "ragged" or "jagged" arrays which include at least one dimension of variable size. The simplest example of this would be an array of variable sized arrays of some simple type like integers or floats.

It may be necessary to divide up this space and split out into separate issues if it would make sense to break the problem down. Suggestions for how to do this welcome.

Some related material:

• The https://github.com/scikit-hep/awkward-array package
• Variable dimensions in ndtypes
• Apache arrow nested data types, including variable sized lists

[alimanfoo]

Just to say that the v3 core protocol currently includes the concepts of data type, array shape, chunk shape, chunk memory layout, array metadata. I imagine that protocol extensions can extend, modify or add to these core protocol features.

Exactly how this works has yet to be fully defined, but there is also some general discussion of protocol extensions in #49.

[jpivarski]

I'll add more later, but I should start with an updated link:

• The Awkward Array 1.x package (different repo from 0.x)

Serialization has been implemented in Awkward 0.x but not yet in Awkward 1.x. This is how it works in the old version:

Given an array with a complex data type,

>>> import awkward as ak0
>>> array = ak0.fromiter([1.1, 2.2, [3, 4, 5], [6, 7], None, {"x": 3.14, "y": 2.72}])
>>> array
<IndexedMaskedArray [1.1 2.2 [3 4 5] [6 7] None <Row 0>] at 0x7fc6f24aa310>
>>> array.tolist()
[1.1, 2.2, [3, 4, 5], [6, 7], None, {'x': 3.14, 'y': 2.72}]

It can be broken down into columns for each node in the data type, and each of these columns is a contiguous, one-dimensional array with a different length.

>>> storage = {}
>>> ak0.persist.serialize(array, storage, name="xyz")
>>> storage.keys()
dict_keys(['xyz-1', 'xyz-4', 'xyz-6', 'xyz-9', 'xyz-10', 'xyz-12', 'xyz-14', 'xyz'])

All of these objects can be stored in any backend that accepts named, binary blobs. One of these is a navigation object (JSON, most likely) and the rest are numerical arrays.

>>> for n, x in storage.items():
...     print(repr(n), repr(x))
... 
'xyz-1' array([ 0,  1,  2,  3, -1,  4])
'xyz-4' array([0, 0, 1, 1, 2], dtype=uint8)
'xyz-6' array([1.1, 2.2])
'xyz-9' array([3, 2])
'xyz-10' array([3, 4, 5, 6, 7])
'xyz-12' array([3.14])
'xyz-14' array([2.72])
'xyz' b'{"awkward": "0.12.20", "schema": {"call": ["awkward", "IndexedMaskedArray"], "args": [{"call": ["awkward", "numpy", "frombuffer"], "args": [{"read": "1"}, {"dtype": "int64"}, {"json": 6, "id": 2}], "id": 1}, {"call": ["awkward", "UnionArray", "fromtags"], "args": [{"call": ["awkward", "numpy", "frombuffer"], "args": [{"read": "4"}, {"dtype": "uint8"}, {"json": 5, "id": 5}], "id": 4}, {"list": [{"call": ["awkward", "numpy", "frombuffer"], "args": [{"read": "6"}, {"dtype": "float64"}, {"json": 2, "id": 7}], "id": 6}, {"call": ["awkward", "JaggedArray", "fromcounts"], "args": [{"call": ["awkward", "numpy", "frombuffer"], "args": [{"read": "9"}, {"dtype": "int64"}, {"ref": 7}], "id": 9}, {"call": ["awkward", "numpy", "frombuffer"], "args": [{"read": "10"}, {"dtype": "int64"}, {"ref": 5}], "id": 10}], "id": 8}, {"call": ["awkward", "Table", "frompairs"], "args": [{"pairs": [["x", {"call": ["awkward", "numpy", "frombuffer"], "args": [{"read": "12"}, {"dtype": "float64"}, {"json": 1, "id": 13}], "id": 12}], ["y", {"call": ["awkward", "numpy", "frombuffer"], "args": [{"read": "14"}, {"dtype": "float64"}, {"ref": 13}], "id": 14}]]}, {"json": 0}], "id": 11}]}], "id": 3}, {"json": -1}], "id": 0}, "prefix": "xyz-"}'

So, for example, these arrays can be stored in HDF5, but only if wrapped with our filter:

>>> h5file = h5py.File("some.hdf5", "a")
>>> ak_h5file = ak0.hdf5(h5file)
>>> ak_h5file
<awkward.hdf5 '/' (0 members)>

>>> ak_h5file["xyz"] = array
>>> ak_h5file["xyz"]
<IndexedMaskedArray [1.1 2.2 [3 4 5] [6 7] None <Row 0>] at 0x7fc6d6c12690>

Here, we interpret the HDF5 group as a single array. If you look at it without the filter, you see all the columns:

>>> h5file["xyz"]
<HDF5 group "/xyz" (8 members)>
>>> h5file["xyz"].keys()
<KeysViewHDF5 ['1', '10', '12', '14', '4', '6', '9', 'schema.json']>

>>> h5file["xyz"]["4"]
<HDF5 dataset "4": shape (5,), type "|u1">
>>> np.asarray(h5file["xyz"]["4"])
array([0, 0, 1, 1, 2], dtype=uint8)

I'm planning to do a similar thing in Awkward 1.x, but in C++ so that it can be used from both C++ and Python. The biggest wart in all of this is that somebody could receive an HDF5 file without being told they need the Awkward wrapper to interpret it, and then it's not easily interpretable.

Since you're developing the v3 extensions interface, I'd like to do this first as a Zarr extension and somehow bake into it the fact that you need the Awkwrad library to interpret it (i.e. no dependency, but some sort of warning or hint of what the user is supposed to do).

The datetime data types extension is similar in one aspect: the underlying data are in primitive types (64-bit integers) while the interpretation is something else (timestamps). However, this extension can be implemented without external libraries and one low-level array is mapped to exactly one high-level array.

For an Awkward extension, a group of different-length low-level arrays are mapped to one high-level array. Zarr has the same concept of a group as HDF5, so it's technically possible to do that mapping, but it would have to be expressed as part of the specification that extensions can map a group to an array.

Could the scope include that?

[jpivarski]

Serialization has been implemented in Awkward 0.x but not yet in Awkward 1.x.

PR scikit-hep/awkward#348 added a similar type of serialization to Awkward 1. There are two new functions, ak.to_arrayset and ak.from_arrayset, which turn an Awkward Array into a schema and a collection of flat arrays (an "arrayset"):

>>> original = ak.Array([[1, 2, 3], [], [4, 5]])
>>> form, container, num_partitions = ak.to_arrayset(original)
>>> form
{
    "class": "ListOffsetArray64",
    "offsets": "i64",
    "content": {
        "class": "NumpyArray",
        "itemsize": 8,
        "format": "l",
        "primitive": "int64",
        "form_key": "node1"
    },
    "form_key": "node0"
}
>>> container
{'node0-offsets': array([0, 3, 3, 5], dtype=int64),
 'node1': array([1, 2, 3, 4, 5])}
>>> print(num_partitions)
None

The flat arrays, 'node0-offsets': array([0, 3, 3, 5]) and 'node1: array([1, 2, 3, 4, 5])' in the above, can be stored in any system that manages named arrays. The extra layer is to group those arrays and pass the metadata (form in the above) to the right functions to reconstruct them (ak.from_arrayset).

If the v3 protocol can make one-to-many associations like this (one user-level object to many physical objects in the "arrayset"), then we should be able to do this.

[jpivarski]

We were just talking about this on the Zarr meeting today (2020-12-16): I've made a minor change to the Awkward <--> group of named buffers transformation. (Motivation and details here: scikit-hep/awkward#592 (comment).) To follow a deprecation procedure properly, the name of the function had to change.

To keep this thread updated, I'll translated the above example into the new interface (awkward>=1.0.1) below.

The following non-rectilinear array,

>>> import awkward as ak
>>> original = ak.Array([[1, 2, 3], [], [4, 5]])
>>> original
<Array [[1, 2, 3], [], [4, 5]] type='3 * var * int64'>

can be decomposed into metadata and flat buffers like this:

>>> mutablemapping = {}
>>> form, length, container = ak.to_buffers(original, container=mutablemapping)

>>> # The values from the original array are now flat buffers in this container, which could be a Zarr sink.
>>> container is mutablemapping
True
>>> container   # dict of NumPy arrays; we only assume they're <memory>
{
    'part0-node0-offsets': array([0, 3, 3, 5], dtype=int64),
    'part0-node1-data': array([1, 2, 3, 4, 5])
}

>>> # The metadata describing how to put those buffers together again is expressed in JSON.
>>> # (It could be included in the container as a byte array.)
>>> form
{
    "class": "ListOffsetArray64",
    "offsets": "i64",
    "content": {
        "class": "NumpyArray",
        "itemsize": 8,
        "format": "l",
        "primitive": "int64",
        "form_key": "node1"
    },
    "form_key": "node0"
}

>>> # We also need to save the length somewhere.
>>> length
3

This goes in the other direction like so:

>>> # Needs the form, length, and container.
>>> ak.from_buffers(form, length, container)
<Array [[1, 2, 3], [], [4, 5]] type='3 * var * int64'>

This page has a full example of saving arbitrary structures to HDF5 using h5py. The problem with that example is that if someone is given an HDF5 file with exploded data in it, they don't know that they need to run ak.from_buffers on it. It will appear to be a bag of one-dimensional arrays not the original array. What I'd want to do with a Zarr v3 extension is to add metadata that says, "I need to import awkward as ak and run ak.from_buffers to extract this group as an array!"

Also in the time since my last comment here, @martindurant has implemented this in Zarr v2: https://github.com/martindurant/awkward_extras/blob/main/awkward_zarr/core.py (though it's based on ak.from_arrayset, not ak.from_buffers, because I only updated that interface earlier this week).

[chrisroat]

Would the main spec be open to allowing arrays where the chunk shapes are not fixed, but the chunk lengths in each dimension are defined by an array (or are constant)?

A concrete example is the rechunking done in dask overlap calculations when the last chunk is smaller than the overlap. A small example is a size 22 1-d array with chunks of 10 -- chunks=(10,10,2). If an overlap of size 3 is needed, dask automatically rechunks to (10,9,3).

Would it be too much overhead to allow each dimension's chunks to be specified as an array? This seems to allow additional flexibility without much extra complexity/space, so could be part of the core spec (instead of as an extension).

[martindurant]

I am also interested in defining chunks as an array along each dimension like dask.array does (I have an issue somewhere) and this would be pretty easy to implement given the existing code in dask that could be copied. It could be done for v2, and leave the existing constant-chunk case unaltered. That would be really useful for cases such as a time dimension wth chunks of 365, 365, 366, 365 days - which you see all the time.

Maybe that's enough for awkward, but I am not sure. Each batch of top-level rows in an awkward array has a variable number of leaf elements within each of the rows. So we can know the total number of leaf values in a batch of rows and call this a "chunk" to be stored together. Maybe that's exactly what we need.

[jpivarski]

In our discussions about chunking Awkward Arrays, we've settled on considering them as logically 1-dimensional, so chunks/partitions would be defined by lengths, not shapes. It wouldn't be a problem to make the length of each chunk/partition constant or given by some list (e.g. 365, 365, 366, 365). I also don't see a problem with the lengths of chunks being 10, 9, 3 instead of 10, 10, 2.

The lengths of internal buffers that constitute those logical arrays are not predictable—or anyway, their lengths depend on or have to agree with values in other arrays. In the above example,

>>> container   # dict of NumPy arrays; we only assume they're <memory>
{
    'part0-node0-offsets': array([0, 3, 3, 5], dtype=int64),
    'part0-node1-data': array([1, 2, 3, 4, 5])
}

the length of partition 0 being 3 means that the buffer at the root of the tree, 'part0-node0-offsets', must have length 4 (fencepost principle for ListOffsetArray), but there's no constraint on the length of 'part0-node1-data' other than that its length must be (at least) equal to part0-node0-offsets[-1], which is 5. With a deeper, more complex tree, there would be a lot of lengths depending on values in their parent nodes' buffers. For a given array, all of the lengths of all of the internal buffers could be specified somewhere, but they would change in ways that depend on values if you ever wanted to change the partitioning.

As an implementation of partition-changing, you could turn a set of buffers for the two partitions that need to be resized into two ak.Arrays, slice them both, and combine the part that needs to be combined with ak.concatenate (in other words, exactly as you would do with NumPy arrays). That would take care of all the internal lengths (ak.from_buffers checks lengths as it descends) and offsets (which have to count from the beginning of a partition, so all of their values increase or decrease by an offset). It necessarily involves accessing all of the data, but anyway, all non-leaf data strictly have to be accessed in order to compute what the new lengths should be. The only possible optimization would be to avoid accessing the buffers for the leaf nodes.

[joshmoore]

My assumption at the current point in the sharding discussion leads me to believe that there should be an extension mechanism (a la codecs) at the array level, e.g. current simple chunks, sharded chunks, etc. Somewhat naively I wonder if this might not be a strategy for irregular chunking but then there's the natural question of how would one then combine sharding and irregular chunking. Thoughts welcome.