Fix: ufuncs on records should not be allowed unless overridden (#1559)

* Added  to recursively_apply for ufuncs on records

* Renamed test to match PR nr

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

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

* Forward behavior in  to check for overriden ufuncs

* Added test for overloaded absolutefunction

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/awkward/_v2/_connect/numpy.py

@@ -271,7 +271,7 @@ def unary_action(layout, **ignore):
                 return result[0]
 
         out = inputs[where].recursively_apply(
-            unary_action, function_name=ufunc.__name__
+            unary_action, behavior, function_name=ufunc.__name__, allow_records=False
         )
 
     else:

src/awkward/_v2/contents/bitmaskedarray.py

@@ -540,7 +540,7 @@ def _completely_flatten(self, nplike, options):
         return self.project()._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if self._nplike.known_shape:
             content = self._content[0 : self._length]
@@ -554,6 +554,7 @@ def continuation():
                     self._mask,
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth,
                         copy.copy(depth_context),
                         lateral_context,
@@ -572,6 +573,7 @@ def continuation():
             def continuation():
                 content._recursively_apply(
                     action,
+                    behavior,
                     depth,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/contents/bytemaskedarray.py

@@ -933,7 +933,7 @@ def _completely_flatten(self, nplike, options):
         return self.project()._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if self._nplike.known_shape:
             content = self._content[0 : self._mask.length]
@@ -947,6 +947,7 @@ def continuation():
                     self._mask,
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth,
                         copy.copy(depth_context),
                         lateral_context,
@@ -963,6 +964,7 @@ def continuation():
             def continuation():
                 content._recursively_apply(
                     action,
+                    behavior,
                     depth,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/contents/content.py

@@ -1319,19 +1319,23 @@ def completely_flatten(self, nplike=None, flatten_records=True, function_name=No
     def recursively_apply(
         self,
         action,
+        behavior=None,
         depth_context=None,
         lateral_context=None,
+        allow_records=True,
         keep_parameters=True,
         numpy_to_regular=True,
         return_array=True,
         function_name=None,
     ):
         return self._recursively_apply(
             action,
+            behavior,
             1,
             depth_context,
             lateral_context,
             {
+                "allow_records": allow_records,
                 "keep_parameters": keep_parameters,
                 "numpy_to_regular": numpy_to_regular,
                 "return_array": return_array,

src/awkward/_v2/contents/emptyarray.py

@@ -313,7 +313,7 @@ def _completely_flatten(self, nplike, options):
         return []
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if options["return_array"]:
 

src/awkward/_v2/contents/indexedarray.py

@@ -1175,7 +1175,7 @@ def _completely_flatten(self, nplike, options):
         return self.project()._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if (
             self._nplike.known_shape
@@ -1196,6 +1196,7 @@ def continuation():
                     index,
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth,
                         copy.copy(depth_context),
                         lateral_context,
@@ -1211,6 +1212,7 @@ def continuation():
             def continuation():
                 content._recursively_apply(
                     action,
+                    behavior,
                     depth,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/contents/indexedoptionarray.py

@@ -1581,7 +1581,7 @@ def _completely_flatten(self, nplike, options):
         return self.project()._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if (
             self._nplike.known_shape
@@ -1606,6 +1606,7 @@ def continuation():
                     index,
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth,
                         copy.copy(depth_context),
                         lateral_context,
@@ -1621,6 +1622,7 @@ def continuation():
             def continuation():
                 content._recursively_apply(
                     action,
+                    behavior,
                     depth,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/contents/listarray.py

@@ -1389,7 +1389,7 @@ def _completely_flatten(self, nplike, options):
             return flat._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if (
             self._nplike.known_shape
@@ -1415,6 +1415,7 @@ def continuation():
                     stops,
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth + 1,
                         copy.copy(depth_context),
                         lateral_context,
@@ -1430,6 +1431,7 @@ def continuation():
             def continuation():
                 content._recursively_apply(
                     action,
+                    behavior,
                     depth + 1,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/contents/listoffsetarray.py

@@ -2027,7 +2027,7 @@ def _completely_flatten(self, nplike, options):
             return flat._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if self._nplike.known_shape and self._nplike.known_data:
             offsetsmin = self._offsets[0]
@@ -2045,6 +2045,7 @@ def continuation():
                     offsets,
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth + 1,
                         copy.copy(depth_context),
                         lateral_context,
@@ -2060,6 +2061,7 @@ def continuation():
             def continuation():
                 content._recursively_apply(
                     action,
+                    behavior,
                     depth + 1,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/contents/numpyarray.py

@@ -1285,11 +1285,11 @@ def _completely_flatten(self, nplike, options):
         return [self.raw(nplike).reshape(-1)]
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if self._data.ndim != 1 and options["numpy_to_regular"]:
             return self.toRegularArray()._recursively_apply(
-                action, depth, depth_context, lateral_context, options
+                action, behavior, depth, depth_context, lateral_context, options
             )
 
         if options["return_array"]:

src/awkward/_v2/contents/recordarray.py

@@ -947,7 +947,7 @@ def _completely_flatten(self, nplike, options):
             )
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if self._nplike.known_shape:
             contents = [x[: self._length] for x in self._contents]
@@ -957,10 +957,17 @@ def _recursively_apply(
         if options["return_array"]:
 
             def continuation():
+                if not options["allow_records"]:
+                    raise ak._v2._util.error(
+                        ValueError(
+                            f"cannot broadcast records in {options['function_name']}"
+                        )
+                    )
                 return RecordArray(
                     [
                         content._recursively_apply(
                             action,
+                            behavior,
                             depth,
                             copy.copy(depth_context),
                             lateral_context,
@@ -981,6 +988,7 @@ def continuation():
                 for content in contents:
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth,
                         copy.copy(depth_context),
                         lateral_context,

src/awkward/_v2/contents/regulararray.py

@@ -1150,7 +1150,7 @@ def _completely_flatten(self, nplike, options):
             return flat._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if self._nplike.known_shape:
             content = self._content[: self._length * self._size]
@@ -1163,6 +1163,7 @@ def continuation():
                 return RegularArray(
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth + 1,
                         copy.copy(depth_context),
                         lateral_context,
@@ -1180,6 +1181,7 @@ def continuation():
             def continuation():
                 content._recursively_apply(
                     action,
+                    behavior,
                     depth + 1,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/contents/unionarray.py

@@ -1318,7 +1318,7 @@ def _completely_flatten(self, nplike, options):
         return out
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if options["return_array"]:
 
@@ -1329,6 +1329,7 @@ def continuation():
                     [
                         content._recursively_apply(
                             action,
+                            behavior,
                             depth,
                             copy.copy(depth_context),
                             lateral_context,
@@ -1347,6 +1348,7 @@ def continuation():
                 for content in self._contents:
                     content._recursively_apply(
                         action,
+                        behavior,
                         depth,
                         copy.copy(depth_context),
                         lateral_context,

src/awkward/_v2/contents/unmaskedarray.py

@@ -502,14 +502,15 @@ def _completely_flatten(self, nplike, options):
         return self.project()._completely_flatten(nplike, options)
 
     def _recursively_apply(
-        self, action, depth, depth_context, lateral_context, options
+        self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if options["return_array"]:
 
             def continuation():
                 return UnmaskedArray(
                     self._content._recursively_apply(
                         action,
+                        behavior,
                         depth,
                         copy.copy(depth_context),
                         lateral_context,
@@ -525,6 +526,7 @@ def continuation():
             def continuation():
                 self._content._recursively_apply(
                     action,
+                    behavior,
                     depth,
                     copy.copy(depth_context),
                     lateral_context,

src/awkward/_v2/highlevel.py

@@ -241,7 +241,9 @@ def __init__(
             )
 
         if with_name is not None:
-            layout = ak._v2.operations.with_name(layout, with_name, highlevel=False)
+            layout = ak._v2.operations.with_name(
+                layout, with_name, highlevel=False, behavior=behavior
+            )
 
         if backend is not None and backend != ak._v2.operations.backend(layout):
             layout = ak._v2.operations.to_backend(layout, backend, highlevel=False)

src/awkward/_v2/operations/ak_cartesian.py

@@ -365,7 +365,9 @@ def getfunction2(layout, depth, **kwargs):
                                 "characters of a string; please split it into lists"
                             )
                         )
-                    nextlayout = layout.recursively_apply(getgetfunction1(inside))
+                    nextlayout = layout.recursively_apply(
+                        getgetfunction1(inside), behavior
+                    )
                     return newaxis(nextlayout, outside)
                 else:
                     return None
@@ -376,7 +378,7 @@ def apply(x, i):
             layout = ak._v2.operations.to_layout(
                 x, allow_record=False, allow_other=False
             )
-            return layout.recursively_apply(getgetfunction2(i))
+            return layout.recursively_apply(getgetfunction2(i), behavior)
 
         toflatten = []
         if nested is None or nested is False:

src/awkward/_v2/operations/ak_fill_none.py

@@ -125,7 +125,7 @@ def action(layout, depth, depth_context, **kwargs):
                 return maybe_fillna(layout)
 
     depth_context = {"posaxis": axis}
-    out = arraylayout.recursively_apply(action, depth_context=depth_context)
+    out = arraylayout.recursively_apply(action, behavior, depth_context=depth_context)
 
     return ak._v2._util.wrap(
         out, ak._v2._util.behavior_of(array, behavior=behavior), highlevel

src/awkward/_v2/operations/ak_from_json.py

@@ -119,7 +119,7 @@ def record_to_complex(node, **kwargs):
     layout = (
         layout
         if complex_imag_string is None
-        else layout.recursively_apply(record_to_complex)
+        else layout.recursively_apply(record_to_complex, behavior)
     )
 
     if highlevel:

src/awkward/_v2/operations/ak_from_json_file.py

@@ -131,7 +131,7 @@ def record_to_complex(node, **kwargs):
     layout = (
         layout
         if complex_imag_string is None
-        else layout.recursively_apply(record_to_complex)
+        else layout.recursively_apply(record_to_complex, behavior)
     )
 
     nonfinite_dict = {}
@@ -147,7 +147,11 @@ def string_to_nonfinite(node, **kwargs):
             if node.parameter("__array__") == "string":
                 return node._awkward_strings_to_nonfinite(nonfinite_dict)
 
-    layout = layout.recursively_apply(string_to_nonfinite) if nonfinite_dict else layout
+    layout = (
+        layout.recursively_apply(string_to_nonfinite, behavior)
+        if nonfinite_dict
+        else layout
+    )
 
     layout = (
         layout.content