Repository: scikit-hep/awkward

Issue:
Allow reducers to be overridden in v2
### Description of new feature

This was the second half of #1375; I'm splitting it because it's really two issues.

Now #1375 is just for applying the policy in v1, which will prevent bugs.

This issue is for extending the capability to what users really want—allowing reducers to be overridden so that `ak.sum` on Vectors will add the Vectors (accounting for coordinate systems). It can only be done in v2.


Discussion:
Motivation:

```python
import vector

v = array of lists of vectors in pt, phi, eta coordinates

ak.sum(v, axis=1)   # CORRECTLY adds vectors in each list to produce one vector per list
```

Vector needs to be able to overload record-sum in such a way as to correctly handle non-Cartesian coordinates. On our side, that means that reducers over records (such as vectors) need to be overloadable.

Overloading syntax?

```python
ak.behaviors[np.sum, "Momentum3D"] = special_function   # implemented in vector, not awkward
```

There's already a syntax for overloading NumPy ufuncs,

```python
ak.behaviors[np.add, "Momentum3D", "Momentum3D"] = special_function
```

Ufuncs are "mapper" functions; the new capability would be to be able to overload "reducer" functions. The only ones that will likely ever be overloaded are `np.sum` (very likely) and `np.prod` (not very likely). Others, like `np.any` or `np.all`, probably not, because they usually act on booleans. Maybe `np.min` and `np.max` would want overloads for records so that orderings can be given to complex objects. (If someone wants to overload `np.min` and `np.max`, they'd want to connect it somehow with the ufuncs `np.less_than`, `np.greater_than`, `np.minimize`, and `np.maximize`...)

Should it just be records, or also arrays? (`__record__: "Momentum3D"` and also `__array__: "string"`?)

Resolving patch:
diff --git a/src/awkward/_do.py b/src/awkward/_do.py
--- a/src/awkward/_do.py
+++ b/src/awkward/_do.py
@@ -204,6 +204,7 @@ def remove_structure(
     function_name: str | None = None,
     drop_nones: bool = True,
     keepdims: bool = False,
+    allow_records: bool = False,
 ):
     if isinstance(layout, Record):
         return remove_structure(
@@ -213,6 +214,7 @@ def remove_structure(
             function_name,
             drop_nones,
             keepdims,
+            allow_records,
         )
 
     else:
@@ -225,6 +227,7 @@ def remove_structure(
                 "function_name": function_name,
                 "drop_nones": drop_nones,
                 "keepdims": keepdims,
+                "allow_records": allow_records,
             },
         )
         return tuple(arrays)
@@ -266,7 +269,11 @@ def reduce(
 ):
     if axis is None:
         parts = remove_structure(
-            layout, flatten_records=False, drop_nones=False, keepdims=keepdims
+            layout,
+            flatten_records=False,
+            drop_nones=False,
+            keepdims=keepdims,
+            allow_records=True,
         )
 
         if len(parts) > 1:
@@ -276,9 +283,9 @@ def reduce(
                 "cannot use axis=None on an array containing irreducible unions"
             )
         elif len(parts) == 0:
-            parts = [ak.contents.EmptyArray()]
-
-        (layout,) = parts
+            layout = ak.contents.EmptyArray()
+        else:
+            (layout,) = parts
 
         starts = ak.index.Index64.zeros(1, layout.backend.index_nplike)
         parents = ak.index.Index64.zeros(layout.length, layout.backend.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
@@ -11,7 +11,7 @@
 from awkward._backends.typetracer import TypeTracerBackend
 from awkward._behavior import find_record_reducer
 from awkward._errors import AxisError
-from awkward._layout import maybe_posaxis
+from awkward._layout import maybe_posaxis, wrap_layout
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpylike import IndexType, NumpyMetadata
 from awkward._nplikes.shape import unknown_length
@@ -35,6 +35,15 @@
 numpy = Numpy.instance()
 
 
+def _apply_record_reducer(
+    reducer, layout: RecordArray, mask: bool, offsets: ak.index.Index, behavior
+) -> Content:
+    # Build a 1D list over these contents
+    array = wrap_layout(ak.contents.ListOffsetArray(offsets, layout), behavior=behavior)
+    # Perform the reduction
+    return ak.to_layout(reducer(array, mask))
+
+
 @final
 class RecordArray(Content):
     """
@@ -889,14 +898,112 @@ def _reduce_next(
         if reducer_recordclass is None:
             raise TypeError(
                 "no ak.{} overloads for custom types: {}".format(
-                    reducer.name, ", ".join(self._fields)
+                    reducer.name, ", ".join(self.fields)
                 )
             )
         else:
-            raise NotImplementedError(
-                "overloading reducers for RecordArrays has not been implemented yet"
+            # Convert parents into offsets
+            outoffsets = ak.index.Index64.empty(
+                outlength + 1, self._backend.index_nplike
+            )
+            assert (
+                outoffsets.nplike is self._backend.index_nplike
+                and parents.nplike is self._backend.index_nplike
+            )
+            self._handle_error(
+                self._backend[
+                    "awkward_ListOffsetArray_reduce_local_outoffsets_64",
+                    outoffsets.dtype.type,
+                    parents.dtype.type,
+                ](
+                    outoffsets.data,
+                    parents.data,
+                    parents.length,
+                    outlength,
+                )
+            )
+
+            out = _apply_record_reducer(
+                reducer_recordclass, self, mask, outoffsets, behavior
             )
 
+            if reducer.needs_position:
+                assert isinstance(out, ak.contents.NumpyArray)
+
+                if shifts is None:
+                    assert (
+                        out.backend is self._backend
+                        and parents.nplike is self._backend.index_nplike
+                        and starts.nplike is self._backend.index_nplike
+                    )
+                    self._handle_error(
+                        self._backend[
+                            "awkward_NumpyArray_reduce_adjust_starts_64",
+                            out.data.dtype.type,
+                            parents.dtype.type,
+                            starts.dtype.type,
+                        ](
+                            out.data,
+                            outlength,
+                            parents.data,
+                            starts.data,
+                        )
+                    )
+                else:
+                    assert (
+                        out.backend is self._backend
+                        and parents.nplike is self._backend.index_nplike
+                        and starts.nplike is self._backend.index_nplike
+                        and shifts.nplike is self._backend.index_nplike
+                    )
+                    self._handle_error(
+                        self._backend[
+                            "awkward_NumpyArray_reduce_adjust_starts_shifts_64",
+                            out.data.dtype.type,
+                            parents.dtype.type,
+                            starts.dtype.type,
+                            shifts.dtype.type,
+                        ](
+                            out.data,
+                            outlength,
+                            parents.data,
+                            starts.data,
+                            shifts.data,
+                        )
+                    )
+
+            if mask:
+                outmask = ak.index.Index8.empty(outlength, self._backend.index_nplike)
+                assert (
+                    outmask.nplike is self._backend.index_nplike
+                    and parents.nplike is self._backend.index_nplike
+                )
+                self._handle_error(
+                    self._backend[
+                        "awkward_NumpyArray_reduce_mask_ByteMaskedArray_64",
+                        outmask.dtype.type,
+                        parents.dtype.type,
+                    ](
+                        outmask.data,
+                        parents.data,
+                        parents.length,
+                        outlength,
+                    )
+                )
+
+                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)
+
+            return out
+
     def _validity_error(self, path):
         for i, cont in enumerate(self.contents):
             if cont.length < self.length:
@@ -1008,6 +1115,8 @@ def _remove_structure(self, backend, options):
             for content in self._contents:
                 out.extend(content[: self._length]._remove_structure(backend, options))
             return out
+        elif options["allow_records"]:
+            return [self]
         else:
             in_function = ""
             if options["function_name"] is not None:
diff --git a/src/awkward/operations/ak_with_name.py b/src/awkward/operations/ak_with_name.py
--- a/src/awkward/operations/ak_with_name.py
+++ b/src/awkward/operations/ak_with_name.py
@@ -44,10 +44,11 @@ def _impl(array, name, highlevel, behavior):
 
     def action(layout, **ignore):
         if isinstance(layout, ak.contents.RecordArray):
-            parameters = dict(layout.parameters)
-            parameters["__record__"] = name
             return ak.contents.RecordArray(
-                layout.contents, layout.fields, layout.length, parameters=parameters
+                layout._contents,
+                layout._fields,
+                layout._length,
+                parameters={**layout.parameters, "__record__": name},
             )
         else:
             return None


Test patch:
diff --git a/tests/test_1607_no_reducers_on_records.py b/tests/test_1607_no_reducers_on_records.py
--- a/tests/test_1607_no_reducers_on_records.py
+++ b/tests/test_1607_no_reducers_on_records.py
@@ -57,11 +57,11 @@ def test_reducers():
 
 
 def test_overloaded_reducers():
-    def overload_add(array, axis=-1):
+    def overload_add(array, mask):
         return ak.contents.RecordArray(
             [
-                ak.contents.NumpyArray(np.asarray([ak.sum(array["rho"], axis=axis)])),
-                ak.contents.NumpyArray(np.asarray([ak.sum(array["phi"], axis=axis)])),
+                ak.sum(array["rho"], axis=-1, mask_identity=False, highlevel=False),
+                ak.sum(array["phi"], axis=-1, mask_identity=False, highlevel=False),
             ],
             ["rho", "phi"],
         )
@@ -75,15 +75,17 @@ def overload_add(array, axis=-1):
         behavior=behavior,
     )
 
-    with pytest.raises(NotImplementedError):
-        assert ak.to_list(ak.sum(array, axis=1)) == [{"rho": 0, "phi": 0}]
+    assert ak.to_list(ak.sum(array, axis=1)) == [{"rho": 0, "phi": 0}]
 
-    with pytest.raises(TypeError):
+    with pytest.raises(
+        TypeError, match=r"overloads for custom types: VectorArray2D, int"
+    ):
         ak.to_list(array + 1)
 
-    def overload_add2(array):
+    def overload_add2(array, mask):
         return ak.contents.RecordArray(
-            [ak.contents.NumpyArray(np.asarray([2.4, 3, 4.5, 6]))], ["rho"]
+            [ak.contents.NumpyArray(np.asarray([2.4, 3, 4.5, 6], dtype=np.float64))],
+            ["rho"],
         )
 
     behavior = {}
@@ -95,8 +97,7 @@ def overload_add2(array):
         behavior=behavior,
     )
 
-    with pytest.raises(NotImplementedError):
-        assert ak.to_list(ak.sum(array, axis=1)) == [{"rho": 2.4}]
+    assert ak.to_list(ak.sum(array, axis=1)) == [{"rho": 2.4}]
 
     array = ak.highlevel.Array(
         ak.contents.ByteMaskedArray(
@@ -115,5 +116,30 @@ def overload_add2(array):
         )
     )
 
-    with pytest.raises(TypeError):
+    with pytest.raises(TypeError, match=r"overloads for custom types: a, b"):
         ak.to_list(ak.sum(array, axis=0))
+
+    def overload_argmax(array, mask):
+        return ak.argmax(array["rho"], axis=-1, mask_identity=False, highlevel=False)
+
+    behavior = {}
+    behavior[ak.argmax, "VectorArray2D"] = overload_argmax
+
+    array = ak.Array(
+        [
+            [
+                {"rho": -1.1, "phi": -0.1},
+                {"rho": 1.1, "phi": 0.1},
+                {"rho": -1.1, "phi": 0.3},
+            ],
+            [
+                {"rho": 1.1, "phi": 0.1},
+                {"rho": -1.1, "phi": 0.3},
+                {"rho": -1.1, "phi": -0.1},
+            ],
+        ],
+        with_name="VectorArray2D",
+        behavior=behavior,
+    )
+
+    assert ak.to_list(ak.argmax(array, axis=1)) == [1, 0]