fix: add highlevel, behavior arguments to composite reducers (#2754)

* fix: add `highlevel`, `behavior` arguments to composite reducers

* test: initial, simple test

* test: check two-arg functions

---------

Co-authored-by: Jim Pivarski <jpivarski@users.noreply.github.com>

src/awkward/_layout.py

@@ -36,6 +36,31 @@ def wrap_layout(content, behavior=None, highlevel=True, like=None, allow_other=F
     return content
 
 
+def maybe_highlevel_to_lowlevel(obj):
+    """
+    Args:
+        obj: an object
+
+    Calls #ak.to_layout and returns the result iff. the object is a high-level
+    Awkward object, otherwise the object is returned as-is.
+
+    This function should be removed once scalars are properly handled by `to_layout`.
+    """
+    import awkward.highlevel
+
+    if isinstance(
+        obj,
+        (
+            awkward.highlevel.Array,
+            awkward.highlevel.Record,
+            awkward.highlevel.ArrayBuilder,
+        ),
+    ):
+        return awkward.to_layout(obj)
+    else:
+        return obj
+
+
 def from_arraylib(array, regulararray, recordarray):
     from awkward.contents import (
         ByteMaskedArray,

src/awkward/operations/ak_corr.py

@@ -3,6 +3,7 @@
 import awkward as ak
 from awkward._behavior import behavior_of
 from awkward._dispatch import high_level_function
+from awkward._layout import maybe_highlevel_to_lowlevel, wrap_layout
 from awkward._nplikes import ufuncs
 from awkward._nplikes.numpylike import NumpyMetadata
 from awkward._regularize import regularize_axis
@@ -11,7 +12,17 @@
 
 
 @high_level_function()
-def corr(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
+def corr(
+    x,
+    y,
+    weight=None,
+    axis=None,
+    *,
+    keepdims=False,
+    mask_identity=False,
+    highlevel=True,
+    behavior=None,
+):
     """
     Args:
         x: One coordinate to use in the correlation (anything #ak.to_layout recognizes).
@@ -33,6 +44,10 @@ def corr(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
             empty lists results in None (an option type); otherwise, the
             calculation is followed through with the reducers' identities,
             usually resulting in floating-point `nan`.
+        highlevel (bool): If True, return an #ak.Array; otherwise, return
+            a low-level #ak.contents.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
 
     Computes the correlation of `x` and `y` (many types supported, including
     all Awkward Arrays and Records, must be broadcastable to each other).
@@ -55,12 +70,12 @@ def corr(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
     yield x, y, weight
 
     # Implementation
-    return _impl(x, y, weight, axis, keepdims, mask_identity)
+    return _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior)
 
 
-def _impl(x, y, weight, axis, keepdims, mask_identity):
+def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior):
     axis = regularize_axis(axis)
-    behavior = behavior_of(x, y, weight)
+    behavior = behavior_of(x, y, weight, behavior=behavior)
     x = ak.highlevel.Array(
         ak.operations.to_layout(x, allow_record=False, allow_other=False),
         behavior=behavior,
@@ -76,8 +91,12 @@ def _impl(x, y, weight, axis, keepdims, mask_identity):
         )
 
     with np.errstate(invalid="ignore", divide="ignore"):
-        xmean = ak.operations.ak_mean._impl(x, weight, axis, False, mask_identity)
-        ymean = ak.operations.ak_mean._impl(y, weight, axis, False, mask_identity)
+        xmean = ak.operations.ak_mean._impl(
+            x, weight, axis, False, mask_identity, highlevel=True, behavior=behavior
+        )
+        ymean = ak.operations.ak_mean._impl(
+            y, weight, axis, False, mask_identity, highlevel=True, behavior=behavior
+        )
         xdiff = x - xmean
         ydiff = y - ymean
         if weight is None:
@@ -130,4 +149,9 @@ def _impl(x, y, weight, axis, keepdims, mask_identity):
                 highlevel=True,
                 behavior=behavior,
             )
-        return sumwxy / ufuncs.sqrt(sumwxx * sumwyy)
+        return wrap_layout(
+            maybe_highlevel_to_lowlevel(sumwxy / ufuncs.sqrt(sumwxx * sumwyy)),
+            behavior=behavior,
+            highlevel=highlevel,
+            allow_other=True,
+        )

src/awkward/operations/ak_covar.py

@@ -3,14 +3,25 @@
 import awkward as ak
 from awkward._behavior import behavior_of
 from awkward._dispatch import high_level_function
+from awkward._layout import maybe_highlevel_to_lowlevel, wrap_layout
 from awkward._nplikes.numpylike import NumpyMetadata
 from awkward._regularize import regularize_axis
 
 np = NumpyMetadata.instance()
 
 
 @high_level_function()
-def covar(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
+def covar(
+    x,
+    y,
+    weight=None,
+    axis=None,
+    *,
+    keepdims=False,
+    mask_identity=False,
+    highlevel=True,
+    behavior=None,
+):
     """
     Args:
         x: One coordinate to use in the covariance calculation (anything #ak.to_layout recognizes).
@@ -32,6 +43,10 @@ def covar(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
             empty lists results in None (an option type); otherwise, the
             calculation is followed through with the reducers' identities,
             usually resulting in floating-point `nan`.
+        highlevel (bool): If True, return an #ak.Array; otherwise, return
+            a low-level #ak.contents.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
 
     Computes the covariance of `x` and `y` (many types supported, including
     all Awkward Arrays and Records, must be broadcastable to each other).
@@ -52,12 +67,12 @@ def covar(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
     yield x, y, weight
 
     # Implementation
-    return _impl(x, y, weight, axis, keepdims, mask_identity)
+    return _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior)
 
 
-def _impl(x, y, weight, axis, keepdims, mask_identity):
+def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior):
     axis = regularize_axis(axis)
-    behavior = behavior_of(x, y, weight)
+    behavior = behavior_of(x, y, weight, behavior=behavior)
     x = ak.highlevel.Array(
         ak.operations.to_layout(x, allow_record=False, allow_other=False),
         behavior=behavior,
@@ -73,8 +88,12 @@ def _impl(x, y, weight, axis, keepdims, mask_identity):
         )
 
     with np.errstate(invalid="ignore", divide="ignore"):
-        xmean = ak.operations.ak_mean._impl(x, weight, axis, False, mask_identity)
-        ymean = ak.operations.ak_mean._impl(y, weight, axis, False, mask_identity)
+        xmean = ak.operations.ak_mean._impl(
+            x, weight, axis, False, mask_identity, highlevel=True, behavior=behavior
+        )
+        ymean = ak.operations.ak_mean._impl(
+            y, weight, axis, False, mask_identity, highlevel=True, behavior=behavior
+        )
         if weight is None:
             sumw = ak.operations.ak_count._impl(
                 x,
@@ -109,4 +128,9 @@ def _impl(x, y, weight, axis, keepdims, mask_identity):
                 highlevel=True,
                 behavior=behavior,
             )
-        return sumwxy / sumw
+        return wrap_layout(
+            maybe_highlevel_to_lowlevel(sumwxy / sumw),
+            behavior=behavior,
+            highlevel=highlevel,
+            allow_other=True,
+        )

src/awkward/operations/ak_linear_fit.py

@@ -15,7 +15,17 @@
 
 
 @high_level_function()
-def linear_fit(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
+def linear_fit(
+    x,
+    y,
+    weight=None,
+    axis=None,
+    *,
+    keepdims=False,
+    mask_identity=False,
+    highlevel=True,
+    behavior=None,
+):
     """
     Args:
         x: One coordinate to use in the linear fit (anything #ak.to_layout recognizes).
@@ -37,6 +47,10 @@ def linear_fit(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=Fa
             empty lists results in None (an option type); otherwise, the
             calculation is followed through with the reducers' identities,
             usually resulting in floating-point `nan`.
+        highlevel (bool): If True, return an #ak.Array; otherwise, return
+            a low-level #ak.contents.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
 
     Computes the linear fit of `y` with respect to `x` (many types supported,
     including all Awkward Arrays and Records, must be broadcastable to each
@@ -72,12 +86,12 @@ def linear_fit(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=Fa
     yield x, y, weight
 
     # Implementation
-    return _impl(x, y, weight, axis, keepdims, mask_identity)
+    return _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior)
 
 
-def _impl(x, y, weight, axis, keepdims, mask_identity):
+def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior):
     axis = regularize_axis(axis)
-    behavior = behavior_of(x, y, weight)
+    behavior = behavior_of(x, y, weight, behavior=behavior)
     backend = backend_of(x, y, weight, coerce_to_common=True, default=cpu)
     x = ak.highlevel.Array(
         ak.operations.to_layout(x, allow_record=False, allow_other=False).to_backend(
@@ -247,4 +261,6 @@ def _impl(x, y, weight, axis, keepdims, mask_identity):
         if scalar:
             out = out[0]
 
-        return wrap_layout(out, highlevel=True, behavior=behavior, allow_other=scalar)
+        return wrap_layout(
+            out, highlevel=highlevel, behavior=behavior, allow_other=scalar
+        )