Repository: scikit-hep/awkward

Issue:
Positional record-reducers cannot return masked results
### Version of Awkward Array

main

### Description and code to reproduce

We provide a `mask` argument that indicates to custom reducers whether a mask should be returned, for performance. For positional reducers, the proper thing would be to set `mask=False` always; we have to do more work for masked results in order to rebuild the index values correctly, and whilst we can use e.g. `fill_none` on the result, the `mask` argument becomes meaningless in such a scenario. Instead, we can set `mask` to `False` for positional reducers, meaning that the user should always return non-option types.


Discussion:

Resolving patch:
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
@@ -929,10 +929,28 @@ def _reduce_next(
                 )
             )
 
+            # 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, self, mask, outoffsets, behavior
+                reducer_recordclass,
+                self,
+                reducer_should_mask,
+                outoffsets,
+                behavior,
             )
 
+            if out.is_option and not reducer_should_mask:
+                reason = (
+                    "reducer is positional"
+                    if reducer.needs_position
+                    else "mask is False"
+                )
+                raise TypeError(
+                    f"a custom implementation of the reducer {reducer.name} for {self.parameter('__record__')!r} "
+                    f"returned an option when it was not expected ({reason})"
+                )
+
             if reducer.needs_position:
                 assert isinstance(out, ak.contents.NumpyArray)
 
@@ -1000,10 +1018,6 @@ def _reduce_next(
                 out = ak.contents.ByteMaskedArray.simplified(
                     outmask, out, False, parameters=None
                 )
-            elif out.is_option:
-                raise TypeError(
-                    "a custom reducer function returned an option when it was not expected"
-                )
 
             if keepdims:
                 out = ak.contents.RegularArray(out, 1, self.length, parameters=None)


Test patch:
diff --git a/tests/test_2501_positional_record_reducer.py b/tests/test_2501_positional_record_reducer.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2501_positional_record_reducer.py
@@ -0,0 +1,151 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import pytest
+
+import awkward as ak
+
+
+def _min_pair(array, mask):
+    array = ak.typetracer.empty_if_typetracer(array)
+
+    # Find location of minimum 0 slot
+    i_min = ak.argmin(array["0"], axis=-1, keepdims=True, mask_identity=True)
+    # Index into array
+    pair_min = ak.ravel(array[ak.from_regular(i_min)], highlevel=False)
+    if mask:
+        return pair_min
+    else:
+        form = pair_min.content.form
+        length_one_content = form.length_one_array(
+            backend=pair_min.backend, highlevel=False
+        )
+        identity_content = ak.fill_none(length_one_content, 0, highlevel=False)
+        identity = ak.record.Record(identity_content, 0)
+        return ak.fill_none(pair_min, identity, highlevel=False)
+
+
+def _argmin_pair(array, mask):
+    array = ak.typetracer.empty_if_typetracer(array)
+
+    assert not mask
+    # Find location of minimum 0 slot
+    return ak.argmin(array["0"], axis=-1, keepdims=False, mask_identity=mask)
+
+
+def _argmin_pair_bad(array, mask):
+    array = ak.typetracer.empty_if_typetracer(array)
+
+    assert not mask
+    # Find location of minimum 0 slot
+    return ak.argmin(array["0"], axis=-1, keepdims=False, mask_identity=True)
+
+
+def test_non_positional():
+    behavior = {(ak.min, "pair"): _min_pair}
+
+    x = ak.Array(
+        [
+            [
+                [1, 2, 3],
+                [5, 4, 3],
+                [2],
+            ],
+            [
+                [8],
+                [],
+                [10, 4, 4],
+            ],
+        ]
+    )
+    y = 2 * x - x**2
+    z = ak.zip((x, y), with_name="pair", behavior=behavior)
+
+    assert ak.almost_equal(
+        ak.min(z, axis=-1, mask_identity=True),
+        ak.Array(
+            [
+                [(1, 1), (3, -3), (2, 0)],
+                [
+                    (8, -48),
+                    None,
+                    (4, -8),
+                ],
+            ],
+            with_name="pair",
+        ),
+    )
+
+    assert ak.almost_equal(
+        ak.min(z, axis=-1, mask_identity=False),
+        ak.Array(
+            [
+                [(1, 1), (3, -3), (2, 0)],
+                [
+                    (8, -48),
+                    (0, 0),
+                    (4, -8),
+                ],
+            ],
+            with_name="pair",
+        ),
+    )
+
+
+def test_positional_bad():
+    behavior = {(ak.argmin, "pair"): _argmin_pair_bad}
+
+    x = ak.Array(
+        [
+            [
+                [1, 2, 3],
+                [5, 4, 3],
+                [2],
+            ],
+            [
+                [8],
+                [],
+                [10, 4, 4],
+            ],
+        ]
+    )
+    y = 2 * x - x**2
+    z = ak.zip((x, y), with_name="pair", behavior=behavior)
+
+    with pytest.raises(TypeError, match=r"'pair' returned an option"):
+        assert ak.almost_equal(
+            ak.argmin(z, axis=-1, mask_identity=True), [[0, 2, 0], [0, None, 1]]
+        )
+
+    with pytest.raises(TypeError, match=r"'pair' returned an option"):
+        assert ak.almost_equal(
+            ak.argmin(z, axis=-1, mask_identity=False), [[0, 2, 0], [0, -1, 1]]
+        )
+
+
+def test_positional_good():
+    behavior = {(ak.argmin, "pair"): _argmin_pair}
+
+    x = ak.Array(
+        [
+            [
+                [1, 2, 3],
+                [5, 4, 3],
+                [2],
+            ],
+            [
+                [8],
+                [],
+                [10, 4, 4],
+            ],
+        ]
+    )
+    y = 2 * x - x**2
+    z = ak.zip((x, y), with_name="pair", behavior=behavior)
+
+    assert ak.almost_equal(
+        ak.argmin(z, axis=-1, mask_identity=True), [[0, 2, 0], [0, None, 1]]
+    )
+
+    assert ak.almost_equal(
+        ak.argmin(z, axis=-1, mask_identity=False), [[0, 2, 0], [0, -1, 1]]
+    )