Repository: scikit-hep/awkward

Issue:
Record reducers need re-ordering
### Version of Awkward Array

main

### Description and code to reproduce

Unlike the NumpyArray reducers, record reducers do not have access to `parents` to ensure the correct ordering of the result. We should carry the result to ensure correct ordering.


Discussion:

Resolving patch:
diff --git a/src/awkward/contents/listoffsetarray.py b/src/awkward/contents/listoffsetarray.py
--- a/src/awkward/contents/listoffsetarray.py
+++ b/src/awkward/contents/listoffsetarray.py
@@ -1577,6 +1577,8 @@ def _reduce_next(
             )
             nextparents = Index64.empty(nextlen, index_nplike)
 
+            # n.b. awkward_ListOffsetArray_reduce_local_nextparents_64 always returns parents that are
+            # monotonically increasing (because it is local)
             assert (
                 nextparents.nplike is index_nplike
                 and self._offsets.nplike is index_nplike
@@ -1704,6 +1706,7 @@ def _rearrange_prepare_next(self, outlength, parents):
                 maxcount,
             )
         )
+
         maxnextparents = index_nplike.index_as_shape_item(_maxnextparents[0])
         nextstarts = Index64.empty(maxnextparents + 1, index_nplike)
         assert nextstarts.nplike is index_nplike and nextparents.nplike is index_nplike
diff --git a/src/awkward/contents/recordarray.py b/src/awkward/contents/recordarray.py
--- a/src/awkward/contents/recordarray.py
+++ b/src/awkward/contents/recordarray.py
@@ -906,35 +906,46 @@ def _reduce_next(
                 )
             )
         else:
+            # Positional reducers ultimately need to do more work when rebuilding the result
+            # so asking for a mask doesn't help us!
+            reducer_should_mask = mask and not reducer.needs_position
+
             # Convert parents into offsets to build a list for axis=1 reduction
             offsets = ak.index.Index64.empty(outlength + 1, self._backend.index_nplike)
             assert (
                 offsets.nplike is self._backend.index_nplike
                 and parents.nplike is self._backend.index_nplike
             )
+            # `parents` are possibly non monotonic increasing, so we must re-order the result
+            # This happens naturally for the `NumpyArray` reducers.
+            carry = ak.index.Index64.empty(outlength, self._backend.index_nplike)
+
+            # Note: if we knew that `negaxis == depth` exclusively for this layout, we could use
+            # the simpler `ListOffsetArray_reduce_local_outoffsets_64`. However, if our parent was reduced,
+            # we would still see `negaxis == depth`, so this kernel has to be used instead.
+            assert carry.nplike is self._backend.index_nplike
             self._backend.maybe_kernel_error(
                 self._backend[
-                    "awkward_ListOffsetArray_reduce_local_outoffsets_64",
+                    "awkward_RecordArray_reduce_nonlocal_outoffsets_64",
                     offsets.dtype.type,
+                    carry.dtype.type,
                     parents.dtype.type,
                 ](
                     offsets.data,
+                    carry.data,
                     parents.data,
                     parents.length,
                     outlength,
                 )
             )
-            layout_to_reduce = ak.contents.ListOffsetArray(offsets, self)
-
-            # Positional reducers ultimately need to do more work when rebuilding the result
-            # so asking for a mask doesn't help us!
-            reducer_should_mask = mask and not reducer.needs_position
             out = _apply_record_reducer(
                 reducer_recordclass,
-                layout_to_reduce,
+                ak.contents.ListOffsetArray(offsets, self),
                 reducer_should_mask,
                 behavior,
             )
+            out = out._carry(carry, allow_lazy=True)
+
             if out.is_option and not reducer_should_mask:
                 reason = (
                     "reducer is positional"


Test patch:
diff --git a/tests/test_2512_record_array_carry.py b/tests/test_2512_record_array_carry.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2512_record_array_carry.py
@@ -0,0 +1,160 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+
+import awkward as ak
+
+
+def _reduce_max_masked(array, mask):
+    assert mask
+    j = ak.from_regular(
+        ak.argmax(array["1"], axis=1, keepdims=True, mask_identity=True)
+    )
+    return ak.flatten(array[j], axis=1)
+
+
+behavior = {}
+behavior[ak.max, "pair"] = _reduce_max_masked
+
+
+def test_axis_0():
+    content = ak.contents.ListArray(
+        ak.index.Index64([0, 2]),
+        ak.index.Index64([2, 4]),
+        ak.contents.ListOffsetArray(
+            ak.index.Index64([0, 3, 6, 9, 11]),
+            ak.contents.RecordArray(
+                [
+                    ak.contents.NumpyArray(np.arange(11, dtype=np.int64)),
+                    ak.contents.NumpyArray(
+                        np.array(
+                            [
+                                0.0,
+                                2.0,
+                                4.0,
+                                6.0,
+                                8.0,
+                                10.0,
+                                1.0,
+                                14.0,
+                                16.0,
+                                18.0,
+                                20.0,
+                            ],
+                            dtype=np.float64,
+                        )
+                    ),
+                ],
+                fields=None,
+                parameters={"__record__": "pair"},
+            ),
+        ),
+    )
+
+    result = ak.max(
+        content,
+        axis=0,
+        keepdims=True,
+        mask_identity=True,
+        behavior=behavior,
+        highlevel=False,
+    )
+
+    expected_result = ak.contents.ListArray(
+        ak.index.Index64([0]),
+        ak.index.Index64([2]),
+        ak.contents.ListArray(
+            ak.index.Index64([0, 3]),
+            ak.index.Index64([3, 6]),
+            ak.contents.IndexedOptionArray(
+                ak.index.Index64([1, 5, 9, 3, 7, 10]),
+                ak.contents.RecordArray(
+                    [
+                        ak.contents.NumpyArray(
+                            np.array([0, 6, 3, 9, 1, 7, 4, 10, 2, 8, 5], dtype=np.int64)
+                        ),
+                        ak.contents.NumpyArray(
+                            np.array(
+                                [0, 1, 6, 18, 2, 14, 8, 20, 4, 16, 10],
+                                dtype=np.float64,
+                            )
+                        ),
+                    ],
+                    fields=None,
+                    parameters={"__record__": "pair"},
+                ),
+            ),
+        ),
+    )
+    assert result.is_equal_to(expected_result)
+
+
+def test_axis_1():
+    content = ak.contents.ListArray(
+        ak.index.Index64([0, 2]),
+        ak.index.Index64([2, 4]),
+        ak.contents.ListOffsetArray(
+            ak.index.Index64([0, 3, 6, 9, 11]),
+            ak.contents.RecordArray(
+                [
+                    ak.contents.NumpyArray(np.arange(11, dtype=np.int64)),
+                    ak.contents.NumpyArray(
+                        np.array(
+                            [
+                                0.0,
+                                2.0,
+                                4.0,
+                                6.0,
+                                8.0,
+                                10.0,
+                                1.0,
+                                14.0,
+                                16.0,
+                                18.0,
+                                20.0,
+                            ],
+                            dtype=np.float64,
+                        )
+                    ),
+                ],
+                fields=None,
+                parameters={"__record__": "pair"},
+            ),
+        ),
+    )
+
+    result = ak.max(
+        content,
+        axis=1,
+        keepdims=True,
+        mask_identity=True,
+        behavior=behavior,
+        highlevel=False,
+    )
+
+    expected_result = ak.contents.RegularArray(
+        ak.contents.ListArray(
+            ak.index.Index64([0, 3]),
+            ak.index.Index64([3, 6]),
+            ak.contents.IndexedOptionArray(
+                ak.index.Index64([1, 5, 9, 3, 7, 10]),
+                ak.contents.RecordArray(
+                    [
+                        ak.contents.NumpyArray(
+                            np.array([0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8], dtype=np.int64)
+                        ),
+                        ak.contents.NumpyArray(
+                            np.array(
+                                [0, 6, 1, 18, 2, 8, 14, 20, 4, 10, 16],
+                                dtype=np.float64,
+                            )
+                        ),
+                    ],
+                    fields=None,
+                    parameters={"__record__": "pair"},
+                ),
+            ),
+        ),
+        size=1,
+    )
+    assert result.is_equal_to(expected_result)