fix: ak.array_equal on unions with reordered index into list child

src/awkward/operations/ak_almost_equal.py

@@ -1,6 +1,7 @@
 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
 
 
+import awkward as ak
 from awkward._backends.dispatch import backend_of
 from awkward._backends.numpy import NumpyBackend
 from awkward._behavior import behavior_of, get_array_class, get_record_class
@@ -113,12 +114,56 @@ def is_approx_dtype(left, right) -> bool:
                     return np.issubdtype(right, family)
         return left == right
 
+    def eager_carry(layout, carry):
+        # Carry eagerly so that same_content_types sees the packed classes. A lazy
+        # carry wraps a RecordArray in an IndexedArray and an UnmaskedArray above
+        # it then becomes an IndexedOptionArray
+        return layout._carry(ak.index.Index64(carry, nplike=backend.nplike), False)
+
+    def project_union(layout, tag):
+        tags = layout.tags.data
+        return eager_carry(
+            layout.contents[tag], layout.index.data[: tags.shape[0]][tags == tag]
+        )
+
+    def project_option(layout):
+        layout = layout.to_IndexedOptionArray64()
+        index = layout.index.data
+        return eager_carry(layout.content, index[index >= 0])
+
+    def to_list_offset_array(layout):
+        if not isinstance(layout, ak.contents.ListArray):
+            return layout.to_ListOffsetArray64(False)
+        offsets = layout._compact_offsets64(True)
+        # repeat needs counts it can cast to intp which is 32-bit on 32-bit platforms
+        counts = backend.nplike.astype(offsets.data[1:] - offsets.data[:-1], np.intp)
+        shifts = backend.nplike.repeat(layout.starts.data - offsets.data[:-1], counts)
+        carry = backend.nplike.arange(offsets[-1], dtype=np.int64) + shifts
+        return ak.contents.ListOffsetArray(
+            offsets, eager_carry(layout.content, carry), parameters=layout.parameters
+        )
+
     def packed_list_content(layout):
-        layout = layout.to_ListOffsetArray64(False)
+        layout = to_list_offset_array(layout)
         return layout.content[layout.offsets[0] : layout.offsets[-1]]
 
+    def packed_node(layout):
+        # Packed layouts have no ListArray and no IndexedArray (except categorical)
+        # but carrying can create them so convert them back
+        if (
+            isinstance(layout, ak.contents.IndexedArray)
+            and layout.parameter("__array__") != "categorical"
+        ):
+            layout = layout.project()
+        if isinstance(layout, ak.contents.ListArray):
+            layout = to_list_offset_array(layout)
+        return layout
+
     def visitor(left, right) -> bool:
-        # Most firstly, check same_content_types before any transformations
+        left = packed_node(left)
+        right = packed_node(right)
+
+        # Most firstly, check same_content_types before any further transformations
         if same_content_types and left.__class__ is not right.__class__:
             return False
 
@@ -256,7 +301,7 @@ def visitor(left, right) -> bool:
         elif left.is_option and right.is_option:
             return backend.nplike.array_equal(
                 left.mask_as_bool(True), right.mask_as_bool(True)
-            ) and visitor(left.project(), right.project())
+            ) and visitor(project_option(left), project_option(right))
         elif left.is_union and right.is_union:
             # After simplification, both unions should have the same number of contents
             if len(left.contents) != len(right.contents):
@@ -293,7 +338,7 @@ def ordered_unique_values(values):
 
             # Now project out the contents, and check for equality
             for i, j in zip(left_tag_order, right_tag_order, strict=True):
-                if not visitor(left.project(i), right.project(j)):
+                if not visitor(project_union(left, i), project_union(right, j)):
                     return False
             return True
 

tests/test_4316_union_equal_reordered_index.py

@@ -0,0 +1,296 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import itertools
+
+import numpy as np
+import pytest
+
+import awkward as ak
+
+
+def test_reordered_index_list_of_records():
+    """Issue #4316, first reproducer: a union whose list-of-records child is reordered."""
+    z = ak.Array([[{"x": 1}], 2, []])
+    p = z[[1, 2, 0]]
+    q = ak.Array(p.tolist())
+
+    assert p.tolist() == q.tolist()
+    assert ak.array_equal(p, q)
+    assert ak.array_equal(q, p)
+    assert ak.array_equal(p, q, same_content_types=False)
+    assert ak.almost_equal(p, q)
+
+
+def _zero_field_union(index, offsets):
+    return ak.Array(
+        ak.contents.UnionArray(
+            ak.index.Index8(np.array([1, 1], dtype=np.int8)),
+            ak.index.Index64(np.array(index, dtype=np.int64)),
+            [
+                ak.contents.RecordArray([], [], 0),
+                ak.contents.ListOffsetArray(
+                    ak.index.Index64(np.array(offsets, dtype=np.int64)),
+                    ak.contents.RecordArray([], [], 1),
+                ),
+            ],
+        )
+    )
+
+
+def test_reordered_index_hand_built_union():
+    """Issue #4316, second reproducer: identical forms, one union index reordered.
+
+    `right` reaches the same values through a permuted union index, so
+    `UnionArray.project` has to reorder the list child's content. The lazy
+    carry that does so turns the child's `RecordArray` into an `IndexedArray`
+    on one side only.
+    """
+    left = _zero_field_union([0, 1], [0, 1, 1])
+    right = _zero_field_union([1, 0], [0, 0, 1])
+
+    assert left.layout.form == right.layout.form
+    assert left.tolist() == [[{}], []]
+    assert right.tolist() == [[{}], []]
+    assert ak.array_equal(left, right)
+    assert ak.array_equal(right, left)
+    assert ak.array_equal(left, right, same_content_types=False)
+    assert ak.almost_equal(left, right)
+
+
+def _unmasked_records_union(index, offsets):
+    return ak.Array(
+        ak.contents.UnionArray(
+            ak.index.Index8(np.array([1, 1, 0], dtype=np.int8)),
+            ak.index.Index64(np.array(index, dtype=np.int64)),
+            [
+                ak.contents.NumpyArray(np.array([2], dtype=np.int64)),
+                ak.contents.ListOffsetArray(
+                    ak.index.Index64(np.array(offsets, dtype=np.int64)),
+                    ak.contents.UnmaskedArray(
+                        ak.contents.RecordArray(
+                            [ak.contents.NumpyArray(np.array([1], dtype=np.int64))],
+                            ["x"],
+                        )
+                    ),
+                ),
+            ],
+        )
+    )
+
+
+def test_reordered_index_unmasked_records():
+    """A list child of nullable records, which is how pyarrow makes a struct.
+
+    Compacting the reordered side's lists carries the records, and a lazy
+    carry would turn the `UnmaskedArray` over them into an
+    `IndexedOptionArray` on that side only.
+    """
+    left = _unmasked_records_union([0, 1, 0], [0, 1, 1])
+    right = _unmasked_records_union([1, 0, 0], [0, 0, 1])
+
+    assert left.layout.form == right.layout.form
+    assert left.tolist() == [[{"x": 1}], [], 2]
+    assert right.tolist() == [[{"x": 1}], [], 2]
+    assert ak.array_equal(left, right)
+    assert ak.array_equal(right, left)
+    assert ak.array_equal(left, right, same_content_types=False)
+    assert ak.almost_equal(left, right)
+
+
+def test_reordered_index_unmasked_records_from_arrow():
+    """The same layout, as `ak.from_arrow` builds it from a dense union."""
+    pa = pytest.importorskip("pyarrow")
+
+    def dense(offsets, lists):
+        return ak.from_arrow(
+            pa.UnionArray.from_dense(
+                pa.array([1, 1, 0], pa.int8()),
+                pa.array(offsets, pa.int32()),
+                [
+                    pa.array([2], pa.int64()),
+                    pa.array(lists, pa.list_(pa.struct([("x", pa.int64())]))),
+                ],
+            )
+        )
+
+    left = dense([0, 1, 0], [[{"x": 1}], []])
+    right = dense([1, 0, 0], [[], [{"x": 1}]])
+
+    assert left.tolist() == right.tolist()
+    assert ak.array_equal(left, right)
+    assert ak.array_equal(right, left)
+    assert ak.almost_equal(left, right)
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        pytest.param([[1, 2], 3, []], id="list-of-numbers"),
+        pytest.param([[{"x": 1}], 2, []], id="list-of-records"),
+        pytest.param([[(1, 2.5)], 2, []], id="list-of-tuples"),
+        pytest.param([[[1, 2]], 2, []], id="list-of-lists"),
+        pytest.param([["one"], 2, []], id="list-of-strings"),
+        pytest.param([[b"one"], 2, []], id="list-of-bytestrings"),
+        pytest.param([[{"x": 1}, None], 2, []], id="list-of-optional-records"),
+        pytest.param([[[{"x": 1}], []], 2, []], id="list-of-lists-of-records"),
+        pytest.param(
+            [[{"y": [1, 2]}, {"y": []}], 2, [{"y": [3]}]],
+            id="list-of-records-of-lists",
+        ),
+        pytest.param([[{"s": "one"}], 2, [{"s": ""}]], id="list-of-records-of-strings"),
+    ],
+)
+@pytest.mark.parametrize("permutation", list(itertools.permutations(range(3))))
+def test_reordered_index_every_permutation(data, permutation):
+    """Reordering a union's index must not change the outcome for any child type."""
+    p = ak.Array(data)[list(permutation)]
+    q = ak.Array(p.tolist())
+
+    assert p.tolist() == q.tolist()
+    assert ak.array_equal(p, q)
+    assert ak.array_equal(q, p)
+    assert ak.array_equal(p, q, same_content_types=False)
+    assert ak.almost_equal(p, q)
+
+
+def test_reordered_index_unequal_arrays():
+    """A reordered union index must not make unequal arrays compare equal."""
+    p = ak.Array([[{"x": 1}], 2, []])[[1, 2, 0]]
+    assert p.tolist() == [2, [], [{"x": 1}]]
+
+    different_value = ak.Array([2, [], [{"x": 999}]])
+    assert not ak.array_equal(p, different_value)
+    assert not ak.array_equal(different_value, p)
+    assert not ak.almost_equal(p, different_value)
+
+    different_lists = ak.Array([2, [{"x": 1}], []])
+    assert not ak.array_equal(p, different_lists)
+    assert not ak.array_equal(different_lists, p)
+    assert not ak.almost_equal(p, different_lists)
+
+
+def _union_of(list_content):
+    return ak.Array(
+        ak.contents.UnionArray(
+            ak.index.Index8(np.array([0, 1, 1], dtype=np.int8)),
+            ak.index.Index64(np.array([0, 1, 0], dtype=np.int64)),
+            [
+                ak.contents.NumpyArray(np.array([2], dtype=np.int64)),
+                list_content,
+            ],
+        )
+    )
+
+
+def test_union_child_list_classes_still_compared():
+    """`same_content_types=True` must still separate regular from variable lists."""
+    var = _union_of(
+        ak.contents.ListOffsetArray(
+            ak.index.Index64(np.array([0, 1, 2], dtype=np.int64)),
+            ak.contents.NumpyArray(np.array([1, 3], dtype=np.int64)),
+        )
+    )
+    regular = _union_of(
+        ak.contents.RegularArray(
+            ak.contents.NumpyArray(np.array([1, 3], dtype=np.int64)), 1
+        )
+    )
+
+    assert var.tolist() == regular.tolist() == [2, [3], [1]]
+    assert not ak.array_equal(var, regular)
+    assert not ak.array_equal(regular, var)
+    assert ak.array_equal(var, regular, check_regular=False)
+
+
+def test_categorical_content_class_still_compared():
+    """Packing keeps a categorical `IndexedArray`, so the class check must too."""
+    categorical = ak.contents.IndexedArray(
+        ak.index.Index64(np.array([0, 1, 0], dtype=np.int64)),
+        ak.contents.NumpyArray(np.array([2, 3], dtype=np.int64)),
+        parameters={"__array__": "categorical"},
+    )
+    plain = ak.contents.NumpyArray(np.array([2, 3, 2], dtype=np.int64))
+
+    assert not ak.array_equal(categorical, plain)
+    assert not ak.array_equal(plain, categorical)
+    assert ak.array_equal(categorical, plain, same_content_types=False)
+    assert ak.array_equal(categorical, categorical)
+
+
+@pytest.mark.parametrize("check_parameters", [True, False])
+def test_categorical_union_child_class_still_compared(check_parameters):
+    """A categorical union child compares as it does outside a union.
+
+    The lazy carry of the plain records used to be an `IndexedArray` too, so
+    `ak.array_equal` returned True here, although it returns False for the
+    same two contents outside a union.
+    """
+    categorical = ak.contents.IndexedArray(
+        ak.index.Index64(np.array([1, 0], dtype=np.int64)),
+        ak.contents.RecordArray(
+            [ak.contents.NumpyArray(np.array([5, 6], dtype=np.int64))], ["y"]
+        ),
+        parameters={"__array__": "categorical"},
+    )
+    plain = ak.contents.RecordArray(
+        [ak.contents.NumpyArray(np.array([6, 5], dtype=np.int64))], ["y"]
+    )
+    left = _union_of(categorical)
+    right = _union_of(plain)
+
+    assert left.tolist() == right.tolist()
+    assert not ak.array_equal(categorical, plain, check_parameters=check_parameters)
+    assert not ak.array_equal(left, right, check_parameters=check_parameters)
+    assert not ak.array_equal(right, left, check_parameters=check_parameters)
+    assert ak.array_equal(
+        left, right, check_parameters=check_parameters, same_content_types=False
+    )
+    assert ak.almost_equal(left, right, check_parameters=check_parameters)
+
+
+def _regular_of_optional_records():
+    records = ak.contents.RecordArray(
+        [ak.contents.NumpyArray(np.array([1, 3], dtype=np.int64))], ["x"]
+    )
+    unmasked = ak.contents.UnmaskedArray(records)
+    indexed = ak.contents.IndexedOptionArray(
+        ak.index.Index64(np.array([0, 1], dtype=np.int64)), records
+    )
+    return (
+        ak.contents.RegularArray(unmasked, 1),
+        ak.contents.RegularArray(indexed, 1),
+    )
+
+
+def test_union_child_option_classes_still_compared():
+    """Projecting a union child must keep an `UnmaskedArray` below it.
+
+    A lazy carry turns it into an `IndexedOptionArray` on both sides, which
+    would make it pass the class check against a real one.
+    """
+    unmasked, indexed = _regular_of_optional_records()
+    left = _union_of(unmasked)
+    right = _union_of(indexed)
+
+    assert left.tolist() == right.tolist()
+    assert not ak.array_equal(unmasked, indexed)
+    assert not ak.array_equal(left, right)
+    assert not ak.array_equal(right, left)
+    assert ak.array_equal(left, right, same_content_types=False)
+    assert ak.almost_equal(left, right)
+
+
+def test_option_content_classes_still_compared():
+    """Projecting an option's content must keep an `UnmaskedArray` below it."""
+    unmasked, indexed = _regular_of_optional_records()
+    mask = ak.index.Index8(np.array([1, 0], dtype=np.int8))
+    left = ak.contents.ByteMaskedArray(mask, unmasked, valid_when=True)
+    right = ak.contents.ByteMaskedArray(mask, indexed, valid_when=True)
+
+    assert left.to_packed().form != right.to_packed().form
+    assert not ak.array_equal(left, right)
+    assert not ak.array_equal(right, left)
+    assert ak.array_equal(left, right, same_content_types=False)
+    assert ak.almost_equal(left, right)