Ragged data in Awkward vs. Scipp

[SimonHeybrock]

Hi!
I am working on Scipp, which supports some types of ragged data (see Scipp Documentation on Binned Data for details).

Given the interest in pydata/xarray#4285 for using Awkward arrays with Xarray, which is conceptually very close to what scipp.DataArray can do, I thought it is high time to kick off a discussion between us. I think there are a lot of similarities but also differences and it would be valuable to understand those, which might at some point even lead to support of a common API or functionality. Please also have a look in https://discuss.scientific-python.org/t/ragged-array-summit/465.

For starters here are high level similarities and differences I am aware of (note that I have used Awkward for little more than 10 minutes):

• Scipp's binned data is conceptually very similar to ak.layout.ListArray64. We store an ND-array of start and stop indices, and a "buffer" (I think Awkward calls this "content")? I managed to convert scipp binned data into an Awkward list array without copying the buffer/content.
• Scipp does not support nesting or similar more complex cases.
• Scipp supports (1-D) scipp.DataArray and scipp.Dataset (similar to xarray.DataArray/Dataset) as the content. This seems similar to using records in Awkward"s content, with slightly more structure. In particular, we distinguish data fields and coordinate fields (and only data fields are operated on in, e.g., arithmetic operations). This does not really matter in terms of data layout and memory handling, i.e., one could imagine swapping out the internals of how things are stored.
• Scipp provides operations for (re)grouping and (re)binning ragged data in multiple dimensions.
• Scipp is multi-threaded out-of-the-box (based on TBB). I could not find info on whether Awkward does this (or whether it relies on and can make use of multi-threading in the underlying (numpy) array library?

[agoose77]

In particular, we distinguish data fields and coordinate fields (and only data fields are operated on in, e.g., arithmetic operations).

How does this behave in practice? Do you throw errors when attempting arithmetic operations on coordinate fields?

Awkward does not have this distinction at the type level, instead we allow users to define how records should be added; we recently changed our model so that records by default do not add together, but library authors can opt-in by writing an implementation. You can look at the vector library as an example of this, although it's a bit indirect as a consequence of supporting multiple "backends" besides Awkward. The design prototype perhaps shows how vector uses Awkward Array in a trivial sense.

I wonder @jpivarski @ioanaif if there are merits for inverting this logic, so that the default case works, and behavior authors can opt-out? I believe one would do this by setting an np.ufunc rule for their behaviors. I don't have a strong opinion yet.

Scipp is multi-threaded out-of-the-box (based on TBB). I could not find info on whether Awkward does this (or whether it relies on and can make use of multi-threading in the underlying (numpy) array library?

Awkward's CPU-kernels are currently single-threaded. We don't use NumPy for many of our "kernels" (the compiled loops that underpin the high-level operations and reductions). We are working towards CUDA support, which would give us local parallelism, although we do not yet have all the requried kernels implemented there (see our roadmap). I'm not totally sure on what kind of work we'll be doing to this in the near future (as in, I genuinely don't know). @jpivarski will have an answer.

Separately, we also have dask-awkward being developed, which will give us distributed computation (and, through multiple-processes, core-scaling). Clearly this is not the same as local parallelism, but it does somewhat address the problem.

[SimonHeybrock]

Fundamentally, Awkward Array needs to understand the type of the data that it operates upon. We can't, for example, operate on arbitrary Python values - we ultimately dispatch to NumPy (or the equivalent array library e.g. CuPy) to compute reductions, etc. At the high level, we restrict the leaf-values of an Array to be

* numeric types (list, int, complex, ...)

* datetime types

* string/bytes types
  and, technically, tuples / records.

I suppose this is my question: Awkward seems to be able to deal with arbitrary arrays records and provides ak.behavior for customizing operations. So what prevents it from being able to deal with other array-like objects (if the user provides custom behaviors)? For example, one could imagine directly using pandas.DataFrame as "content" (without representing it as a dict of NumPy arrays). Right now Awkward seems to only accept subclasses of Content, so pandas.DataFrame cannot by used directly, but why is that? Is there a fundamental reason?

[agoose77]

I won't speak for @jpivarski, who might have a different take on things (and ultimately steers where this project is going), so I'll just put my short, incomplete take here.

At a design-level, Awkward Array is built around the concept of columnar arrays. This is both for performance and for ergonomics, as Awkward was originally intended for use in HEP, where (among a great deal of other things) it is important to be able to read subsets of large datasets by column. So, were we to support non-Numpy leaves, we'd want them to be structureless. The most appropriate generalisation on Awkward's existing behaviour would be to support non-Numpy flat array-like objects (with no "fields") at the leaves. I don't think this would be impossible (but I'll wait for @jpivarski to inevitably correct me ;)). However, outside of the ufunc mechanism (for which we delegate to NumPy), most structure / reduction operators involve calling a compiled kernel. Many kernels don't care about the underlying content, but the reduction kernels do. Specifically, the reduction kernels generalise on the traditional reducers from NumPy by reducing across sublists, e.g. whilst NumPy can only handle

>>> np.min([1, 2, 3, 4])
1

Awkward supports variable-length sublists, roughly:

>>> min([
    [1, 2, 3, 4],
    [5, 6, 7]
])
[1, 5]

Reductions along any axis ultimately invoke the compiled reduction kernels with a 2D array of variable-length sublists. So, were we to generalise the NumpyArray content to support this, there would need to be some kind of mechanism to implement and define these kernels. If these implementations were written in Python, Awkward would be a significant factor slower.

At a technical level, we assume that the internal Awkward content nodes form a closed set, so we don't have any support for subclassing or modifying these content classes.

We are exploring efforts to go the other way - i.e. Awkward Arrays as Pandas extension types: https://github.com/scikit-hep/awkward/wiki#pandas-integration

[jpivarski]

I think it would be reasonable for to_layout to recognize objects that have an __array__ method and call that to get a NumPy array, which it can then wrap in the Content subclass named NumpyArray. This should be near the end of the list of types to_layout checks, so that (for instance) CuPy arrays don't get copied over to the CPU to turn them into NumPy arrays, which is what happens when you call their __array__ method. We should think about the consequences of this, because there are a lot of array-like types out there for which calling __array__ is destructive, such as Pint arrays with units: would users rather get an error when the units are stripped away, at least to know that it's happening?

Then, for instance, a pandas.Series would get turned into an np.ndarray that can be a leaf-node in an Awkward Array's Content tree. But again, that's information-losing: it removes the Series's index. Maybe we should think about that more before just doing it.

Right now Awkward seems to only accept subclasses of Content, so pandas.DataFrame cannot by used directly, but why is that? Is there a fundamental reason?

So yes, that's fundamental. Data in an Awkward Array are not managed by Python, and it's so that we can do operations like the reducer @agoose77 described. Also some advanced slices and restructuring operations.

There is a little flexibility in what type of object holds the numerical data, wrapped by the NumpyArray class. Right now, that object can be a NumPy ndarray, a CuPy ndarray, or a JAX DeviceArray (for autodiff). Whenever we add a new backend, we also define a new "nplike" that is a shim to make all of these array libraries duck-typable—out of the box, they're maybe 90% compatible. Then we also have to make sure that our own functions that go beyond the rectilinear array libraries, such as @agoose77's reducer example, can be executed on the backends' data. (Hence, we need to duplicate our suite of CPU functions to also work on GPU pointers, which is not done yet.) So the set of backends is slowly and laboriously extended; it's not something in which arbitrary objects can be dropped in. If they were arbitrary objects, there would be a lot of slicing operations, reducers, etc. that we wouldn't be able to implement (without doing them in slow Python for loops).

[SimonHeybrock]

We should think about the consequences of this, because there are a lot of array-like types out there for which calling __array__ is destructive, such as Pint arrays with units: would users rather get an error when the units are stripped away, at least to know that it's happening?

There is a little flexibility in what type of object holds the numerical data, wrapped by the NumpyArray class. Right now, that object can be a NumPy ndarray, a CuPy ndarray, or a JAX DeviceArray (for autodiff). Whenever we add a new backend, we also define a new "nplike" that is a shim to make all of these array libraries duck-typable—out of the box, they're maybe 90% compatible.

This sounds like you are almost there? If you can (1) work on arbitrary NumPy-like buffers, (2) have a method to request a list/dict of columns from the content (similar to __array__ you mention, but supporting record-like objects such as pandas.DataFrame or scipp.DataArray, i.e., returning multiple columns/buffers), and (3) have a way of creating new content (a callable provided by the content?)?

Probably a constructor callable is not sufficient for (3), since in binary operations will also need to operate on the non-column content such as the unit in you pint example. In Scipp we have a similar situation: All kernels are applied to array elements as well as the array unit.

[jpivarski]

We can operate on arbitrary NumPy-like buffers, so it's not hard to write (for instance) an ak.from_dataframe function to zero-copy wrap a Pandas DataFrame. (It would put each Series.values in a NumpyArray, all the columns into a RecordArray, and wrap that up as an ak.Array, maybe carry some metadata as parameters.) But in such a case, we are controlling what we want to wrap and how we wrap it; I'm wavering on the idea of opening up the floodgates and ingesting anything with an __array__ method.

Similarly for scipp.DataArray. We have a lot of functions named like ak.from_* and ak.to_*. That would put the conversion in Awkward, though by symmetry, it could also go in Scipp. (I don't have a strong opinion about that. We consider the set of from/to functions to be an open-ended set and can be third-party library-specific.) Zero-copy and round-trip fidelity are goals, but for technical reasons, it might be necessary to give up on one or the other.