ak.array_equal to override np.array_equal

[martindurant]

Description of new feature

It would be nice to have convenience functions for asserting the equivalence of ak arrays.

We can compare an array's:

• values
• type (i.e., datashape description)
• form (including lengths of sub structures, if known)

Currently, tests tend to compare via .tolist() (resulting in python lists and dicts) which can compare the array values. The result of ak.type is already comparable. There are, however, well-defined rules of whether an array can be considered equal as a whole to another array.

Note that this is not the same as __eq__, which is an element-wise operation.

[jpivarski]

And one more level: layouts can differ even when the Form is known (a content can be longer than than it needs to be, but such "hidden" elements are not considered to impact the equality of two arrays).

For example, these are to be considered the same, even though they have different NumpyArrays (the first counts up to 5 and the second counts up to 10):

>>> ak.Array(ak.layout.ListOffsetArray64(
...     ak.layout.Index64(np.array([0, 3, 3, 5])),
...     ak.layout.NumpyArray(np.arange(5) * 1.1)))
<Array [[0, 1.1, 2.2], [], [3.3, 4.4]] type='3 * var * float64'>

>>> ak.Array(ak.layout.ListOffsetArray64(
...     ak.layout.Index64(np.array([0, 3, 3, 5])),
...     ak.layout.NumpyArray(np.arange(10) * 1.1)))
<Array [[0, 1.1, 2.2], [], [3.3, 4.4]] type='3 * var * float64'>

Both of these arrays have the same Type (3 * var * float64) and the same Form ({"class": "ListOffsetArray64", "offsets": "i64", "content": "float64"}), but they have different buffers. If you ran them through ak.to_buffers and compared the results, those results would be different from each other.

The "lists of rules" of what is considered a valid array are given in the documentation of each Content subclass, in the constructor of each Pythonic illustration. These rules are also checked for a given array in the implementation of ak.is_valid. While there are these explicit lists of what is valid, there aren't any explicit lists of what is equal, though two arrays that have the same .tolist() representation and the same type and are both valid are equal to each other.

[ivirshup]

If I were to call ak.to_buffers(ak.to_packed(awk)) and compare the results of those for equality, are there any big edge cases left beyond Union[Option[X], Y] != Union[Option[X], Option[Y]] (https://github.com/scikit-hep/awkward/issues/2182#issuecomment-1412318839)?

E.g.:

import awkward as ak, numpy as np


def equal_awkward(a: ak.Array, b: ak.Array) -> bool:
    if len(a) != len(b):
        return False

    a_form, _, a_buffers = ak.to_buffers(ak.to_packed(a))
    b_form, _, b_buffers = ak.to_buffers(ak.to_packed(b))

    if a_form != b_form:
        return False
    if set(a_buffers.keys()) != set(b_buffers.keys()):
        return False
    for k in a_buffers.keys():
        if not np.array_equal(a_buffers[k], b_buffers[k], equal_nan=True):
            return False    

    return True

[agoose77]

It depends upon how you define equality. If you mean at the ak.type(array) + buffers level, then packed doesn't produce a totally unique layout. We have a helper function internally to address that.

For example, an IndexedArray can wrap a RecordArray, but this is not visible at the type level. Additionally, there are different option types that are type-level equivalent, but form-level distinct. Finally, there is RegularArray and 2D+ NumpyArray that are both regular. That's all off the top of my head.

[ivirshup]

It depends upon how you define equality.

I think I'd say "interchangeable at the high level API". So I could pass either to computations and get the same result (within floating point accuracy).

The function you point out looks great for this use-case. Are there barriers to exporting it?

[jpivarski]

y = ak.packed(x)

should satisfy:

• x.tolist() == y.tolist()
• x.type == y.type (or str(x.type) == str(y.type); there are different levels on which types can be compared, but and I think that the equivalence classes of Type.__eq__ are the same as those of Type.__str__)

I'm also pretty sure that ak._util.arrays_approx_equal checks the equivalent of the above (faster than actually calling tolist, though).

The thing that ak.packed changes is the layout/form.

The function you point out looks great for this use-case. Are there barriers to exporting it?

NumPy has a np.testing submodule. Maybe we should have an ak.testing submodule and put this there?

[agoose77]

Ah, I have confused two separate thought processes when writing my reply, which is something picked up on. Indeed, the ak.type matches. I meant to say that it's not trivial to get the type and projected buffers of the layouts, which is what we want for fast comparison!

I agree that it might make sense. I think dask-awkward roll something similar for their test suite.

[ivirshup]

NumPy has a np.testing submodule. Maybe we should have an ak.testing submodule and put this there?

I typically expect functions under a .testing module to throw exceptions when things aren't equal, and to trade efficiency for helpful reporting in those exceptions. I think this is closer to np.allclose, np.array_equal, pd.DataFrame.equals. My preference would be for something in the main namespace.

---

Also, if I wanted to use this function right now: how safe do you think calling ak._utils.arrays_approx_equal is? Should I vendor it?

[jpivarski]

For it to go in the main namespace (as an "operation", through the ak.operations submodule that gets promoted to the main namespace), its interface and name should match a NumPy function if it can. What we ought to do is provide overrides ak.allclose and ak.array_equal, but this function also checks type, which those don't. We'd have to think about how to extend ak.allclose and ak.array_equal into a world that includes more types than shape and dtype.

You're right that all of the functions in np.testing are assert_this and assert_that, whereas arrays_approx_equal is not. We shouldn't start down this path of divergence from NumPy, so I'm going to put PR #2198 in draft mode until we know what to do about it.

If you access ak._utils.arrays_approx_equal temporary, it will work until we find a permanent home for it. If you have any released versions of AnnData that depend on that and we move it, then there will be version combinations of Awkward and AnnData that don't work together, which can't be predicted in advance. (We don't know exactly which version will have a moved, public arrays_approx_equal, or if we'll change its interface to make it a better match to a chosen NumPy function.) But that may be okay: you might just say that the first few versions of AnnData with Awkward support will be quickly superseded and only support the latest release. I can say with confidence that ak._utils.arrays_approx_equal won't move or change except for making it public and with a more NumPy-like name and interface, and you'll be in the loop for that.

[agoose77]

You're right that all of the functions in np.testing are assert_this and assert_that, whereas arrays_approx_equal is not. We shouldn't start down this path of divergence from NumPy, so I'm going to put PR https://github.com/scikit-hep/awkward/pull/2198 in draft mode until we know what to do about it.

Personally I think this is fine. Our function is specialised to Awkward-specific features. The NumPy testing submodule is primarily for supporting NumPy's test suite AFAICT. The package doesn't provide any overload hooks, for example. So I think we can re-use the module name in Awkward.

As for the other functions, I think only array_equal is most obvious; allclose would probably need to just invoke ak.all(ak.isclose(...)), which is a ufunc. Therefore, we'd want something more general to be a separate function.

[jpivarski]

np.isclose isn't a ufunc.

>>> type(np.isclose)
<class 'function'>

But woah—it just works! How did that happen?

>>> np.isclose(
...     ak.Array([[1, 2, 3], [], [4, 5]]),
...     ak.Array([[1.00001, 2.00001, 3.00001], [], [4.00001, 5.00001]])
... )
<Array [[True, True, True], [], [True, True]] type='3 * var * bool'>

[jpivarski]

This conversation is happening in two places; let's all move to #2198 for further discussion of what should happen to arrays_approx_equal and isclose.

[agoose77]

I'm going to re-open this, as we technically don't have array_equal implemented here.