Should ak.Array have a .var accessor for jagged indexing?

[agoose77]

TL;DR

1. Awkward adds an extra "advanced" indexing routine that shadows the existing NumPy advanced indexing.
2. It is hard to write code that uses this without adding multiple ak.from_regular calls to ensure the index is jagged.
3. To fix this, we need to make the advanced Awkward-like (jagged) indexing more explicit (or drop NumPy-like indexing altogether)

Motivation

NumPy Compatability

Awkward Array supports both ak.Arrays and np.ndarrays in its high-level API, and in many of the array/layout methods. To maintain compatibility with NumPy, parts of the Awkward API that overload the NumPy API (e.g. ak.sum) behave differently according to whether they are given "NumPy-like" (regular) arrays or jagged arrays. There are several places that this happens:

• Array indexing (ak.Array.__getitem__)
• Array broadcasting (ak.broadcast_arrays)

Indexing mechanisms (array[...] subscript operator)

ak.Array.__getitem__ supports both the NumPy advanced indexing mechanism, and an Awkward-specific one. Although not strictly part of the NumPy API (It isn't a ufunc, or overloaded by __array_function__), the behaviour of the array subscript operator is a constrained by the social contract that Arrays behave like ndarrays. So, to support both kinds of indexing, we switch between the implementations according to the kind of array that is passed into the index, i.e. whether it is entirely jagged, or entirely regular.

Awkward-like Indexing

>>> x = ak.Array( [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]])
>>> ix = ak.Array([[0, 1], [0, 1], [0, 1], [0, 2]]) 
>>> x[ix].tolist()
[[0, 1], [3, 4], [6, 7], [9, 11]]

vs NumPy-like Indexing

>>> x[ak.to_regular(ix)].tolist()
[[[0, 1, 2], [3, 4, 5]],
 [[0, 1, 2], [3, 4, 5]],
 [[0, 1, 2], [3, 4, 5]],
 [[0, 1, 2], [6, 7, 8]]]

I will subsequently refer to indexing with jagged (var dimension) arrays as "awkward-like indexing", whilst indexing with regular arrays is "NumPy-like indexing".

Whilst it is reasonable for Awkward Arrays to behave differently to NumPy ndarrays, a regular Awkward (non-var) Array is conceptually a (stricter) sub-type of a jagged Awkward (var) Array. By a hand-wavy application of Liskov's substitution principle, it should be possible to substitute the former for the latter. However, because we use jaggedness to request Awkward-like semantics, code that interacts with regular arrays nearly always does not work for irregular ones.

Intent vs Structure

Part of the problem here is that we mix the intent of the operation (e.g. perform advanced Awkward indexing) with the structure (e.g. this array has a dimension of fixed size). Whilst this satisfies some notion of compatibility (indexing with NumPy-like arrays should behave like NumPy), it makes it harder to reason about what code will do without actually running it.

For example, consider using argmax of one array to index another. Here is the first code one might write:

def y_by_max_x(x, y):
    ix = ak.argmax(x, axis=-1, keepdims=True)
    return y[ix]

Now, from the rules of indexing, we cannot know if this will succeed for any arrays x and y, because if x is regular, then keepdims=True will give ix a constant * 1 dimension. If x is jagged in the final dimension, ix will be given * var in the last dimension. Not only that, we need to be careful that the leading dimensions of x are all jagged, because otherwise the index will still fail. Thus, the "safe" code becomes

def y_by_max_x(x, y):
    ix = ak.argmax(x, axis=-1, keepdims=True)
    for i in range(1, ix.ndim):
        ix = ak.from_regular(ix, axis=i)
    return y[ix]

This is a bit of a contrived example, but it reflects the fact that just by wanting to index using Awkward-like semantics, we have to add boilerplate to assert our intent (ahead of time).

Proposals

The take-away of the above section is that I would like to separate the different kinds of advanced indexing such that there cannot be any ambiguity. In an ideal world, we wouldn't change the meaning of an operation according to the array dimension types at all, but as aforementioned, we do this for compatibility.

Making a change to solve this will allow the reader to reason about the code, and remove the need for so many guards. @jpivarski and I have had a few conversations on this already, and the following ideas were discussed. For completeness, I include ideas that aren't really strongly supported!

1. Add ak.Array.var that "only" performs Awkward-like indexing:
   • Should this only accept arrays (i.e. not slices, integer indices etc)?
   • Established convention e.g. .loc in Pandas
   • Should we remove this behaviour from ak.Array.__getitem__?
2. Change the indexing type on a per-array basis, e.g. y = ak.with_awkward_semantics(y) such that y[ix] always performs Awkward-like indexing. This would propagate through most operations (but what happens with ak.zip?)
   • This could be done by setting a custom behaviour that calls .var in (1)!
3. Add a special index object that signals intent, e.g. y[ak.jagged[ix]] or y[ix.jagged]
   • This is quite numpy-like (see https://numpy.org/neps/nep-0021-advanced-indexing.html proposal)

Pinging @nsmith whom @jpivarski mentioned had thoughts on an ak.Array.pick-like function in the past

[agoose77]

So, onto discussion. My opinions:

• (1) is perhaps the easiest to implement, and is very explicit, but:
  • If we want to "solve" the question of ambiguity, one would nearly always need to use .var (which is a lot of extra typing, and visual noise) and, drop jagged indexing from [...], otherwise the same problem exists but in reverse.
• (2) allows the user to make a one-off decision, but
  • it disables NumPy-like advanced indexing, which is probably useful (albeit less often than Awkward-like).
  • it still requires the user to defensively call with_awkward_semantics in examples like the one above
  • it has a non-local effect (e.g. indexing operation at line 8 depends on something at line 0).
• (3) has interesting aesthetics, (there's something nice about having indexing and masking being treated as different characterisations of the same operation) but there is precedent to make this a method (2) because we already have ak.Array.mask