Fixing #1266 (in v1 and v2), possibly by reordering nextparents. (#1274)

* Reordered nextparents to be monotonic (in v2), but this isn't right.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [skip-ci] add a 4d study case

* Fix: use mergesort, which is stable for like elements

* This fixes the typetracer failures.

* Fix: use `stable` instead of `mergesort`

`stable` is the proper option for this requirement.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ianna Osborne <ianna.osborne@cern.ch>
Co-authored-by: Angus Hollands <goosey15@gmail.com>

src/awkward/_v2/_typetracer.py

@@ -549,9 +549,9 @@ def searchsorted(self, *args, **kwargs):
         # haystack, needle, side="right"
         raise NotImplementedError
 
-    def argsort(self, *args, **kwargs):
+    def argsort(self, array, *args, **kwargs):
         # array
-        raise NotImplementedError
+        return TypeTracerArray(np.int64, array.shape)
 
     ############################ manipulation
 

src/awkward/_v2/contents/listoffsetarray.py

@@ -1656,14 +1656,12 @@ def _reduce_next(
             )
 
             distincts_length = outlength * maxcount[0]
-            nextcarry = ak._v2.index.Index64.empty(nextlen, self._nplike)
-            nextparents = ak._v2.index.Index64.empty(nextlen, self._nplike)
+            np_nextcarry = self._nplike.empty(nextlen, dtype=np.int64)
+            np_nextparents = self._nplike.empty(nextlen, dtype=np.int64)
             maxnextparents = ak._v2.index.Index64.empty(1, self._nplike)
             distincts = ak._v2.index.Index64.empty(distincts_length, self._nplike)
             assert (
-                nextcarry.nplike is self._nplike
-                and nextparents.nplike is self._nplike
-                and maxnextparents.nplike is self._nplike
+                maxnextparents.nplike is self._nplike
                 and distincts.nplike is self._nplike
                 and self._offsets.nplike is self._nplike
                 and offsetscopy.nplike is self._nplike
@@ -1672,16 +1670,16 @@ def _reduce_next(
             self._handle_error(
                 self._nplike[
                     "awkward_ListOffsetArray_reduce_nonlocal_preparenext_64",
-                    nextcarry.dtype.type,
-                    nextparents.dtype.type,
+                    np_nextcarry.dtype.type,
+                    np_nextparents.dtype.type,
                     maxnextparents.dtype.type,
                     distincts.dtype.type,
                     self._offsets.dtype.type,
                     offsetscopy.dtype.type,
                     parents.dtype.type,
                 ](
-                    nextcarry.data,
-                    nextparents.data,
+                    np_nextcarry,
+                    np_nextparents,
                     nextlen,
                     maxnextparents.data,
                     distincts.data,
@@ -1694,6 +1692,11 @@ def _reduce_next(
                 )
             )
 
+            # A "stable" sort is essential for the subsequent steps.
+            reorder = self._nplike.argsort(np_nextparents, kind="stable")
+            nextcarry = ak._v2.index.Index64(np_nextcarry[reorder])
+            nextparents = ak._v2.index.Index64(np_nextparents[reorder])
+
             nextstarts = ak._v2.index.Index64.empty(maxnextparents[0] + 1, self._nplike)
             assert (
                 nextstarts.nplike is self._nplike and nextparents.nplike is self._nplike

studies/reducers.py

@@ -1506,3 +1506,78 @@ def IndexedOptionArray_reduce_next(self, negaxis, parents, length):
     [101*31, 103*37, 107*41, 109*43, 113*47],
     [     1],
     [  53*2,   59*3,   61*5,   67*7,  71*11]]
+
+content = RawArray([1,2,3, 4,5,6, 5,6,7, 8,9,10, 9,10,11, 12,13,14, 13,14,15, 16,17,18])
+assert content.tolist() == [1,2,3, 4,5,6, 5,6,7, 8,9,10, 9,10,11, 12,13,14, 13,14,15, 16,17,18]
+
+depth1 = RegularArray(content, 3)
+depth2 = RegularArray(depth1, 2)
+depth3 = RegularArray(depth2, 2)
+depth4 = RegularArray(depth3, 2)
+print(list(depth4))
+assert list(depth4) == [[[[[1, 2, 3],[4, 5, 6]],[[5, 6, 7],[8, 9, 10]]],[[[9, 10, 11],[12, 13, 14]],[[13, 14, 15],[16, 17, 18]]]]]
+print("depth1", list(depth1.reduce(-1)))
+print("depth2", list(depth2.reduce(-1)))
+print("depth3", list(depth3.reduce(-1)))
+print("depth4 (-1)", list(depth4.reduce(-1)))
+# >>> np.prod(array, axis=-1)
+# array([[[   6,  120],
+#         [ 210,  720]],
+#
+#        [[ 990, 2184],
+#         [2730, 4896]]])
+assert list(depth4.reduce(-1)) == [[[[6,120],[210,720]],[[990,2184],[2730,4896]]]]
+
+print("depth4 (-2)", list(depth4.reduce(-2)))
+# >>> np.prod(array, axis=-2)
+# array([[[  4,  10,  18],
+#         [ 40,  54,  70]],
+#
+#        [[108, 130, 154],
+#         [208, 238, 270]]])
+assert list(depth4.reduce(-2)) == [[[[4,10,18],[40,54,70]],[[108,130,154],[208,238,270]]]]
+
+print("depth4 (-3)", list(depth4.reduce(-3)))
+# >>> np.prod(array, axis=-3)
+# array([[[  5,  12,  21],
+#         [ 32,  45,  60]],
+#
+#        [[117, 140, 165],
+#         [192, 221, 252]]])
+assert list(depth4.reduce(-3)) == [[[[5,12,21],[32,45,60]],[[117,140,165],[192,221,252]]]]
+
+print("depth4 (-4)", list(depth4.reduce(-4)))
+# >>> np.prod(array, axis=-4)
+# array([[[  9,  20,  33],
+#         [ 48,  65,  84]],
+#
+#        [[ 65,  84, 105],
+#         [128, 153, 180]]])
+assert list(depth4.reduce(-4)) == [[[[6,120],[210,720]],[[990,2184],[2730,4896]]]]
+
+
+
+# >>> np.prod(array, axis=0)
+# array([[[  9,  20,  33],
+#         [ 48,  65,  84]],
+#
+#        [[ 65,  84, 105],
+#         [128, 153, 180]]])
+# >>> np.prod(array, axis=1)
+# array([[[  5,  12,  21],
+#         [ 32,  45,  60]],
+#
+#        [[117, 140, 165],
+#         [192, 221, 252]]])
+# >>> np.prod(array, axis=2)
+# array([[[  4,  10,  18],
+#         [ 40,  54,  70]],
+#
+#        [[108, 130, 154],
+#         [208, 238, 270]]])
+# >>> np.prod(array, axis=3)
+# array([[[   6,  120],
+#         [ 210,  720]],
+#
+#        [[ 990, 2184],
+#         [2730, 4896]]])

tests/v2/test_1271-fix-4D-reducers.py

@@ -0,0 +1,17 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import pytest  # noqa: F401
+import numpy as np  # noqa: F401
+import awkward as ak  # noqa: F401
+
+
+def test():
+    akarray = ak._v2.highlevel.Array(
+        [[[[1], [4]], [[5], [8]]], [[[9], [12]], [[13], [16]]]]
+    )
+    nparray = np.array([[[[1], [4]], [[5], [8]]], [[[9], [12]], [[13], [16]]]])
+
+    assert ak._v2.sum(akarray, axis=3).tolist() == np.sum(nparray, axis=3).tolist()
+    assert ak._v2.sum(akarray, axis=2).tolist() == np.sum(nparray, axis=2).tolist()
+    assert ak._v2.sum(akarray, axis=1).tolist() == np.sum(nparray, axis=1).tolist()
+    assert ak._v2.sum(akarray, axis=0).tolist() == np.sum(nparray, axis=0).tolist()