Ragged DataArrays (Variables) and coordinates

[jpivarski]

From a conversation with @TomNicholas, @dcherian, and others at lunch today (and, ultimately, following up on #4285): we could perhaps get ragged arrays into xarray by

1. Limiting Awkward Array to only variable- and fixed-length lists of numbers:

import awkward as ak

def is_ragged_type(t):       # formality  # variable-length        # fixed-length
    if isinstance(t, (ak.types.ArrayType, ak.types.ListType, ak.types.RegularType)):
        return is_ragged_type(t.content)
    elif isinstance(t, ak.types.NumpyType):
        return True
    else:
        return False

some_data = ak.Array([[[1.1, 2.2, 3.3], [], [4.4, 5.5]], [], [[6.6, 7.7, 8.8, 9.9]]])
assert is_ragged_type(some_data.type)

2. Allowing shape to return type tuple[int | None] (safer than interpreting -1 or some other number to mean "variable."
3. Instead of coordinates being $N$ one-dimensional arrays, the coordinates for an $N$ dimensional ragged array would be $N$ increasingly ragged arrays.

To explain (3) above, here's an example of coordinates that work with some_data:

x = ak.Array([                    100, 200,            300])
y = ak.Array([[       10, 20,    30 ],  [], [          40]])
z = ak.Array([[[1, 2, 3], [], [4, 5]],  [], [[6, 7, 8, 9]]])

The first dimension, x, looks familiar; it's a one-dimensional array with the same length as some_data. The second dimension, y, is an array of lists in which each list has the same length as the first-level lists in some_data. The third dimension, z, has the same-length lists of lists as some_data.

This structure allows for alignment between coordinates and the data for the same reason that Awkward broadcasting works, demonstrated with + below:

>>> # just showing them for comparison
>>> x.show(type=True)
type: 3 * int64
[100,
 200,
 300]
>>> y.show(type=True)
type: 3 * var * int64
[[10, 20, 30],
 [],
 [40]]
>>> z.show(type=True)
type: 3 * var * var * int64
[[[1, 2, 3], [], [4, 5]],
 [],
 [[6, 7, 8, 9]]]
>>> some_data.show(type=True)
type: 3 * var * var * float64
[[[1.1, 2.2, 3.3], [], [4.4, 5.5]],
 [],
 [[6.6, 7.7, 8.8, 9.9]]]

>>> # broadcast some_data and x
>>> (some_data + x).show(type=True)
type: 3 * var * var * float64
[[[101, 102, 103], [], [104, 106]],
 [],
 [[307, 308, 309, 310]]]

>>> # broadcast some_data and y
>>> (some_data + y).show(type=True)
type: 3 * var * var * float64
[[[11.1, 12.2, 13.3], [], [34.4, 35.5]],
 [],
 [[46.6, 47.7, 48.8, 49.9]]]

>>> # broadcast some_data and z
>>> (some_data + z).show(type=True)
type: 3 * var * var * float64
[[[2.1, 4.2, 6.3], [], [8.4, 10.5]],
 [],
 [[12.6, 14.7, 16.8, 18.9]]]

This isn't a direct extension of aligning rectilinear arrays, but 1D array × 1D array × 1D array is inherently rectilinear. Whereas in rectilinear alignment, we might have a 1D array z like

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

this alignment is a generalization of using a z like

>>> np.lib.stride_tricks.as_strided(z, (3, 3, 3), (0, 0, z.itemsize))
array([[[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]],

       [[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]],

       [[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]]])

@dcherian thought that xarray handles things in a sufficiently abstract way that we might be able to introduce this as a generalization. I'd like to see how that would go (and can provide ideas, implementations, and features on the Awkward end to help).

Also, go ahead and ping anyone here!

[TomNicholas]

Thanks for raising this (and so quickly) @jpivarski. I will read through it very slowly later on, feel free to ping me if I forget to respond later.

cc'ing @jbusecke @mgrover1 as fellow scipy lunch participants, and those who indicated interest in the previous thread: @joshmoore, @milancurcic, @SimonHeybrock, + @shoyer for opinions.

[dcherian]

@jpivarski Can you share some code to define .shape?

[jpivarski]

Yes, I'll try to do this later, but for now, I can say that it will be the recursive function shown above, yielding an integer length from ArrayType, a numerical size from RegularType, and the non-numerical None token from variable-length ListType (nothing from the NumpyArray).

[dcherian]

No rush. I parsed typestr to get shape and dtype in the meantime 🤣

[TomNicholas]

I parsed typestr to get shape and dtype in the meantime rofl

Haha I did exactly that in an experiment a while ago

[jpivarski]

As long as it passes is_ragged_type, parsing str(some_data.type) would work fine. I'll still get you a better function, though.

(Just landed, after a flight...)

[jpivarski]

def shape(t):
    if isinstance(t, ak.types.ArrayType):
        return (t.length,) + shape(t.content)
    elif isinstance(t, ak.types.RegularType):
        return (t.size,) + shape(t.content)
    elif isinstance(t, ak.types.ListType):
        return (None,) + shape(t.content)
    elif isinstance(t, ak.types.NumpyType):
        return ()
    else:
        raise TypeError("oops, is_ragged_type would have returned False")

and then

>>> shape(some_data.type)
(3, None, None)

and

>>> regular_data = ak.Array(np.arange(2*3*5).reshape(2, 3, 5))
>>> shape(regular_data.type)
(2, 3, 5)

[Illviljan]

Should it be None for undefined? Dask uses np.nan for undefined values for example.

[shoyer]

None is a nice sentinel value, but I am concerned that it could be misinterpretted in Xarray. I can imagine many types of non-integer arrays sizes:

• Ragged dimensions with an array (via Awkward Array)
• Inconsistent sizes between arrays on the same Dataset (each of which could still be NumPy arrays)
• Unknown array sizes, but which are guaranteed to be consistent between arrays on the same Dataset (e.g., unknown shapes in TensorFlow or Dask)

So if we do this, I would lean towards using a more explicit system of some sort, e.g., a xarray.RAGGED sentinel object.

[SimonHeybrock]

Not suggesting it (since I don't know your requirements), but maybe it is useful to compare to a similar but different approach, taken by Scipp, illustrated below:

• Only the regular dimensions of the ragged array contribute the dims and shape of the Variable.
• Ragged coordinates are stored as part of the (contiguous) content buffer, nested in the Variable, alongside the ragged data.
• Since the dims and shape are benign (no None), using this Variable as the "data" in a DataArray is trivial.
• Coords of the "outer" DataArray may be associated with those nested inside the "data" Variable. This is the result of, e.g., non-destructively binning data in one or more dimensions. Note that the association is by name only, i.e., there is no explicit code/link in the data structures.

Binned:

Grouped:

[shoyer]

This version seems like it would be relatively easier to fit into Xarray's existing data model.

[swamidass]

It would be great to see support for ragged arrays. We use them all the time for ML on graphs.

Wouldn't an easier approach be to build a ragged array class that is backed by multiple non-ragged xarrays? This is the approach that, for example, Tensorflow uses. RaggedTensors are composite arrays, each of which is backed by more than one normal tensor.

Perhaps even a ragged "view" of all the ragged tensors derived dataset could be constructed automatically if the right metadata is stored in the attrs.

By building on top of xarray this way, you'd avoid the challenge of navigating all the xarray internals.

[jpivarski]

Can xarray-datatree have different-length DataArrays so that it can encode list content and list lengths (as @swamidass is describing)? I saw in your talk that it addresses ocean gliders with variable sampling rates, which is one of the use-cases I was looking at earlier, so it sounds like xarray-datatree covers it.

[TomNicholas]

Each node of a datatree effectively stores a distinct xarray Dataset, each of which can have any length dimensions you want. Multi-resolution data is one of the main use cases of datatree.

[swamidass]

@TomNicholas not quite. I'll explain later. That's really not a workable approach to this particular problem.

@jpivarski Im happy to explain more the graph use case and why ak package doesn't work so well. Basically there are two key needs:

1. A good data loading and management library that:

• is numpy restricted. other small dependencies are fine, but nothing large or with portability issues, or interferes with major tensor libraries (onnx, torch, Jax, tf).

• keeps all the meta data together in a clean way.

• supports lazy loading of files.

As you can see, Xarray fits the bill here quite nicely, even including dask as a way to parallelize. It is missing ragged integration (Datatree does not work), but that's workable with composite tensors (I'll add explain in another response).

2. A backend agnostic library for native ragged array calculations (the key two basic ops are segment reduce and broadcast, from which most everything can be constructed). Key backends to include are tensorflow, numpy (strict!), Jax and torch. This library would need to provide a consistent set of ops with identical API for all four backends, and ideally also a context manager for switching between backends in different parts of the code (usually between numpy and one of tensorback ends).

Right now this library just doesn't exist. Leaving aside ragged arrays, this doesn't exist for even just normal array ops.

If it did exist, though, and ak built itself around such a library, it might be perfect for us.

This would enable us to write complicated functions, as complex as the leading graph neural networks in one function. One function, based on context, could work natively with any one of the tensor packages, and also operate in the non-GPU preprocessing and deployment contexts (which are a many!).

If Ak allowed us to do that, we switch over to it tomorrow. Is that something you'd be interested in?

If so, I do think I recently found a surprising way to accomplish the main blocker: building that backend agnostic tensor library. If that worked, I'd be curious if ak could produce low dependency (and perhaps reduced feature) API that restricts itself to this set of ops.

If you did that would be a compelling reason to move over. Though some attention would have to be put into ensuring some key details that are critical for interoperating with tensor libraries and our graph use case.

[swamidass]

Thanks for the feed back everyone.

@TomNicholas Datatrees aren't the best way tto do this. Yes, everything has a different size, but (1) not invariant constraints are enforced across all datatrees, and (2) this will scale horribly inefficiently with the number of elements in each ragged dimension.

Remember that Dataset's can have arrays of differnet sizes too, but in that context at least soem of the necessary constraints can be enforced. I'd represnet a ragged array of shape (N, ragged, M) like this:

ds["values"] = xr.DataArray(values, dims=("R", "M"))
ds["row_lengths"] = xr.DataArray(row_lengths, dims=("N",))

The row lengths are meant to define an ragged partition of the R dimension, and it's an open quesiton the best way to record that in the metadata. The good things about this representation is that concatenation is easy.

Alternatively, you could represent it like this:

seg_ids = np.repeat(np.arange(len(row_lengths)), row_lengths)
ds["values"] = xr.DataArray(values, dims=("R", "M"))
ds.coord["R"] = xr.DataArray(seg_ids, dims=("N",))

The nices thing about this representation is that some indexing/selection operations would be a lot easer.

[jpivarski]

I responded to @swamidass's comments/questions about Awkward Array on scikit-hep/awkward#2814.

This thread can be about getting ragged array features in the xarray interface, which is not what Awkward Array is. (And I see value in having something be a true drop-in replacement for xarray, but with ragged arrays.)

[dcherian]

The most recent Array Cast Episode seems to describe a RaggedArray type structure with named axes, implemented in Mathematica. https://www.arraycast.com/episodes/episode66-tali-beynon

https://math.tali.link/rainbow-array-algebra/

[jpivarski]

Follow up on this topic at scikit-hep/ragged#6

[TomNicholas]

There's a massive discussion thread about how one might do this in zarr here: zarr-developers/zarr-specs#62, but I'm not sure if any concrete suggestions exist there.

[dcherian]

@ocefpaf do you know of any CF-compliant netCDF or CDL files for these DSG datasets? http://cfconventions.org/Data/cf-conventions/cf-conventions-1.10/cf-conventions.html#representations-features

[dcherian]

Ah NCEI has templates: https://www.ncei.noaa.gov/data/oceans/ncei/formats/netcdf/v2.0/index.html (of course!) @jpivarski if you scroll down you should see template files with dummy data for the "orthogonal" and "incomplete" representations.

[jpivarski]

Thanks! I'll take a look soon.

[jpivarski]

I was just going through these documents today and I don't see any examples of "contiguous" or "indexed" (i.e. ragged) representations. There are files in the table associated with CDL Template "orthogonal" (totally regular) and "incomplete" (padded to be regular).

Only one .nc file on this page looks like data: stearns_wharf-2013.nc, in a row of the table labeled "orthogonal". Internally, all of the datasets are either scalars (lat, lon)have length 25 (station) or length 16976 (everything else). It doesn't even seem to be multidimensional.

I suppose, taking this as an example, that the connected instance dimensions would be connected by a naming convention, such as

• pressure
• pressure_flagPrimary
• pressure_flagSecondary

and similarly for temperature, salinity, conductivity, and chlorophyll in this example file. Since this is orthogonal, there's no example of what to do with the element dimension, either a "count variable" or an "index variable." If these things are connected only by naming conventions, you probably need the ragged library's interface to be loose and opinionated, to let users align names however they need them to be aligned.

In that case, the interface in #7988 (reply in thread) might be all that you want. Is it? If so, I'll add a test to pin that interface and merge the PR so that someone can try it out.

[dcherian]

array_ufunc and array_function could be used to allow xarray to wrap duck arrays, but since ragged already has array_namespace it wouldn't be used either way.

Seems like we should fallback to the older protocols when we can. Especially since this will get us nan* methods.

[Zac-HD]

I could help with that. In unit tests, I generally make ragged arrays from Python nested lists.
The trickier thing is to make some dimensions regular. We can fix a size in the list comprehension, but upon import it would be interpreted as a None-sized dimension

I would do this differently. If you look at the new documentation on our xr.Variable-generating hypothesis strategies, you'll see that generally we generate the array shape first, then populate it with data afterwards. I'm talking about generalising strategies.array_shapes from returning SearchStrategy[Tuple[int, ...]], to SearchStrategy[Tuple[int | None, ...]], with the actual data values being generated by a secondary array_strategy_fn which accepts an array shape (see #6903 (files)).

I'd probably generate a rectilinear array first, then convert to ragged, and finally truncate each sublist at a chosen length or to the last non-fillvalue. This makes the interpretation of the underlying choices somwhat more stable, which in turn improves performance in generating variations on known examples - and shrinking as a special case of that. We use a similar trick inside mutually_broadcastable_shapes().

---

In fact it seems odd that hypothesis's built-in array_shapes strategy for the array API appears not to generate None-sized dimensions.

Maybe this is because when you decide that a shape item is None, you're not done defining the size of the container: you have to generate random lengths for the lists in that dimension.

More precisely, it's because "an array dimension must be None if and only if a dimension is unknown", and the creation functions do not accept dimensions of size None. Since the array_shapes() strategy primarily aims to help create arrays, we don't support None - but you could write st.lists(st.integers(min_side, max_side) | st.none(), min_size=min_dims, max_size=max_dims).map(tuple) if you wanted to.

scikit-hep/ragged#6 says:

The Array API's purpose in allowing None as a shape item is to accommodate dimensions of unknown length, rather than variable length: not the same thing! This library, Ragged, is an experiment in reinterpreting that part of the standard to see what happens. If it works, that is, if Ragged easily composes with other libraries that consume the standard, then perhaps future versions of the standard can include language that explicitly allows it [, or provide an alternative].

If a future version of the Array API standard includes ragged arrays, we will of course aim to support it in Hypothesis. Until then, I think strategies for generating ragged arrays would be best-placed in the libraries defining those structures, or a third-party package (e.g. hypothesis-ragged) if upstream would prefer it split out. I remain very happy to consult on designs and/or review code for anyone working on such a project.

[jpivarski]

I'd probably generate a rectilinear array first, then convert to ragged, and finally truncate each sublist at a chosen length or to the last non-fillvalue.

That also works, and it can also be done efficiently. The example implementation I wrote above uses ListOffsetArray to select ranges of an underlying NumpyArray to be lists using offsets, integers like 0, 3, 5, 10 to represent a length-3, length-2, length-5 sequence of lists, which necessarily select abutting ranges of the underlying array.

What you're describing could be done with ListArray to select ranges using separate "starts" and "stops" indexes, integers like 0, 3, 5 and 3, 5, 10 (to use the same example), but they don't need to be contiguous. That way, we could start with a rectilinear array and chop off bits of each equal-sized list to make them unequal.

However, the implementation above avoids any for-loops over the scale of the data (only matters for large datasets) and is deterministic (the property I think you're more interested in). The same inputs would always produce the same arrays.

the creation functions do not accept dimensions of size None.

That's right, and it's why this was built using functions beyond the Array API. The Array API doesn't define an array constructor, and ragged.array's constructor adopts the ak.Array that was built manually.

Since a shape with None would be uninformative in a function like ragged.zeros, these functions are equivalent to NumPy's and can only build rectilinear arrays. Functions like ragged.zeros_like can build ragged arrays by adopting the ragged shape from another array. There's no way within the Array API to define a specific ragged shape, specifying all of the list lengths.

[Zac-HD]

the implementation above avoids any for-loops over the scale of the data (only matters for large datasets) and is deterministic (the property I think you're more interested in).

Hypothesis is generally for small (if tricky) data, and determinism is important - but the property I was going for here is locality of changes. That is, if Hypothesis tries to make similar underlying choices, we want our 'converter' code to produce a similar final value for us to test.

Using ListOffsetArray, if we choosing a different offset for one list that will shift the contents of every following list (or affect lengths) - whereas if we're choosing our stop indices along a rectilinear array, our variations will affect each sublist in isolation. That's not a major concern, to be clear, but when you start composing multiple arrays and constrains on them the small performance impacts multiply together pretty quickly.

---

More generally, I just wouldn't expect Hypothesis' array-api support to be particularly helpful for ragged arrays, as more than a starting point and inspiration. I'd be happy to help with review etc. for a (e.g.) ragged.testing.strategies module though, and @honno might also be interested in helping out.