fix: add return_value='simplified' to ak.transform and revamp ak.firs…

…ts/ak.singletons (#1968)

* fix: hide _recursively_apply, add return_value='simplified' to ak.transform, and revamp ak.firsts

* Revamped ak.singletons and added axis=XYZ to tests.

* Introduced Content.is_leaf and corrected some posaxis checking.

* style: pre-commit fixes

* Explain the axis=0 special case with a comment.

Co-authored-by: Angus Hollands <[email protected]>

* Revert 'recursively_apply' hiding (TBD in another PR).

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

src/awkward/contents/bitmaskedarray.py

@@ -628,9 +628,13 @@ def _recursively_apply(
             content = self._content
 
         if options["return_array"]:
+            if options["return_simplified"]:
+                make = BitMaskedArray.simplified
+            else:
+                make = BitMaskedArray
 
             def continuation():
-                return BitMaskedArray(
+                return make(
                     self._mask,
                     content._recursively_apply(
                         action,

src/awkward/contents/bytemaskedarray.py

@@ -1011,9 +1011,13 @@ def _recursively_apply(
             content = self._content
 
         if options["return_array"]:
+            if options["return_simplified"]:
+                make = ByteMaskedArray.simplified
+            else:
+                make = ByteMaskedArray
 
             def continuation():
-                return ByteMaskedArray(
+                return make(
                     self._mask,
                     content._recursively_apply(
                         action,

src/awkward/contents/content.py

@@ -40,6 +40,7 @@ class Content:
     is_indexed = False
     is_record = False
     is_union = False
+    is_leaf = False
 
     def _init(self, parameters: dict[str, Any] | None, backend: Backend):
         if parameters is None:
@@ -1531,6 +1532,7 @@ def recursively_apply(
         allow_records: bool = True,
         keep_parameters: bool = True,
         numpy_to_regular: bool = True,
+        return_simplified: bool = True,
         return_array: bool = True,
         function_name: str | None = None,
     ) -> Content | None:
@@ -1544,6 +1546,7 @@ def recursively_apply(
                 "allow_records": allow_records,
                 "keep_parameters": keep_parameters,
                 "numpy_to_regular": numpy_to_regular,
+                "return_simplified": return_simplified,
                 "return_array": return_array,
                 "function_name": function_name,
             },

src/awkward/contents/emptyarray.py

@@ -15,6 +15,7 @@
 
 class EmptyArray(Content):
     is_unknown = True
+    is_leaf = True
 
     def __init__(self, *, parameters=None, backend=None):
         if backend is None:
@@ -299,7 +300,7 @@ def _pad_none(self, target, axis, depth, clip):
         posaxis = self.axis_wrap_if_negative(axis)
         if posaxis != depth:
             raise ak._errors.wrap_error(
-                np.AxisError(f"axis={axis} exceeds the depth of this array({depth})")
+                np.AxisError(f"axis={axis} exceeds the depth of this array ({depth})")
             )
         else:
             return self.pad_none_axis0(target, True)

src/awkward/contents/indexedarray.py

@@ -1030,9 +1030,13 @@ def _recursively_apply(
             index, content = self._index, self._content
 
         if options["return_array"]:
+            if options["return_simplified"]:
+                make = IndexedArray.simplified
+            else:
+                make = IndexedArray
 
             def continuation():
-                return IndexedArray(
+                return make(
                     index,
                     content._recursively_apply(
                         action,

src/awkward/contents/indexedoptionarray.py

@@ -1561,9 +1561,13 @@ def _recursively_apply(
             index, content = self._index, self._content
 
         if options["return_array"]:
+            if options["return_simplified"]:
+                make = IndexedOptionArray.simplified
+            else:
+                make = IndexedOptionArray
 
             def continuation():
-                return IndexedOptionArray(
+                return make(
                     index,
                     content._recursively_apply(
                         action,

src/awkward/contents/numpyarray.py

@@ -16,6 +16,7 @@
 
 class NumpyArray(Content):
     is_numpy = True
+    is_leaf = True
 
     def __init__(self, data, *, parameters=None, backend=None):
         if backend is None:
@@ -384,7 +385,9 @@ def _offsets_and_flattened(self, axis, depth):
             return self.to_RegularArray()._offsets_and_flattened(posaxis, depth)
 
         else:
-            raise ak._errors.wrap_error(np.AxisError("axis out of range for flatten"))
+            raise ak._errors.wrap_error(
+                np.AxisError(f"axis={axis} exceeds the depth of this array ({depth})")
+            )
 
     def _mergeable(self, other, mergebool):
         if isinstance(
@@ -497,7 +500,7 @@ def _local_index(self, axis, depth):
             return self._local_index_axis0()
         elif len(self.shape) <= 1:
             raise ak._errors.wrap_error(
-                np.AxisError("'axis' out of range for local_index")
+                np.AxisError(f"axis={axis} exceeds the depth of this array ({depth})")
             )
         else:
             return self.to_RegularArray()._local_index(posaxis, depth)
@@ -1083,7 +1086,7 @@ def _combinations(self, n, replacement, recordlookup, parameters, axis, depth):
             return self._combinations_axis0(n, replacement, recordlookup, parameters)
         elif len(self.shape) <= 1:
             raise ak._errors.wrap_error(
-                np.AxisError("'axis' out of range for combinations")
+                np.AxisError(f"axis={axis} exceeds the depth of this array ({depth})")
             )
         else:
             return self.to_RegularArray()._combinations(
@@ -1227,7 +1230,7 @@ def _pad_none(self, target, axis, depth, clip):
         posaxis = self.axis_wrap_if_negative(axis)
         if posaxis != depth:
             raise ak._errors.wrap_error(
-                np.AxisError(f"axis={axis} exceeds the depth of this array({depth})")
+                np.AxisError(f"axis={axis} exceeds the depth of this array ({depth})")
             )
         if not clip:
             if target < self.length:

src/awkward/contents/recordarray.py

@@ -19,6 +19,10 @@
 class RecordArray(Content):
     is_record = True
 
+    @property
+    def is_leaf(self):
+        return len(self._contents) == 0
+
     def __init__(
         self,
         contents,

src/awkward/contents/regulararray.py

@@ -105,6 +105,18 @@ def __deepcopy__(self, memo):
     def simplified(cls, content, size, zeros_length=0, *, parameters=None):
         return cls(content, size, zeros_length, parameters=parameters)
 
+    @property
+    def offsets(self):
+        return self._compact_offsets64(True)
+
+    @property
+    def starts(self):
+        return self._compact_offsets64(True)[:-1]
+
+    @property
+    def stops(self):
+        return self._compact_offsets64(True)[1:]
+
     def _form_with_key(self, getkey):
         form_key = getkey(self)
         return self.Form(

src/awkward/contents/unionarray.py

@@ -1513,9 +1513,13 @@ def _recursively_apply(
         self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if options["return_array"]:
+            if options["return_simplified"]:
+                make = UnionArray.simplified
+            else:
+                make = UnionArray
 
             def continuation():
-                return UnionArray.simplified(
+                return make(
                     self._tags,
                     self._index,
                     [

src/awkward/contents/unmaskedarray.py

@@ -468,9 +468,13 @@ def _recursively_apply(
         self, action, behavior, depth, depth_context, lateral_context, options
     ):
         if options["return_array"]:
+            if options["return_simplified"]:
+                make = UnmaskedArray.simplified
+            else:
+                make = UnmaskedArray
 
             def continuation():
-                return UnmaskedArray(
+                return make(
                     self._content._recursively_apply(
                         action,
                         behavior,

src/awkward/operations/ak_firsts.py

@@ -44,24 +44,58 @@ def firsts(array, axis=1, *, highlevel=True, behavior=None):
 
 
 def _impl(array, axis, highlevel, behavior):
-    layout = ak.operations.to_layout(array, allow_record=False, allow_other=False)
+    layout = ak.operations.to_layout(array)
+    behavior = ak._util.behavior_of(array, behavior=behavior)
     posaxis = layout.axis_wrap_if_negative(axis)
 
+    if not ak._util.is_integer(axis):
+        raise ak._errors.wrap_error(
+            TypeError(f"'axis' must be an integer, not {axis!r}")
+        )
+
     if posaxis == 0:
-        if len(layout) == 0:
-            out = None
+        # specialized logic; it's tested in test_0582-propagate-context-in-broadcast_and_apply.py
+        # Build an integer-typed slice array, so that we can
+        # ensure we have advanced indexing for both length==0
+        # and length > 0 cases.
+        slicer = ak.from_iter([None, 0])
+        if layout.length == 0:
+            out = layout[slicer[[0]]][0]
         else:
-            out = layout[0]
+            out = layout[slicer[[1]]][0]
+
     else:
-        if posaxis < 0:
-            raise ak._errors.wrap_error(
-                NotImplementedError("ak.firsts with ambiguous negative axis")
-            )
-        toslice = (slice(None, None, None),) * posaxis + (0,)
-        out = ak.operations.mask(
-            layout,
-            ak.operations.num(layout, axis=posaxis) > 0,
-            highlevel=False,
-        )[toslice]
-
-    return ak._util.wrap(out, behavior, highlevel, like=array)
+
+        def action(layout, depth, depth_context, **kwargs):
+            posaxis = layout.axis_wrap_if_negative(depth_context["posaxis"])
+            if posaxis == depth and layout.is_list:
+                nplike = layout._backend.index_nplike
+
+                # this is a copy of the raw array
+                index = starts = nplike.array(layout.starts.raw(nplike), dtype=np.int64)
+
+                # this might be a view
+                stops = layout.stops.raw(nplike)
+
+                empties = starts == stops
+                index[empties] = -1
+
+                return ak.contents.IndexedOptionArray.simplified(
+                    ak.index.Index64(index), layout._content
+                )
+
+            elif layout.is_leaf:
+                raise ak._errors.wrap_error(
+                    np.AxisError(
+                        f"axis={axis} exceeds the depth of this array ({depth})"
+                    )
+                )
+
+            depth_context["posaxis"] = posaxis
+
+        depth_context = {"posaxis": posaxis}
+        out = layout.recursively_apply(
+            action, behavior, depth_context, numpy_to_regular=True
+        )
+
+    return ak._util.wrap(out, behavior, highlevel)

src/awkward/operations/ak_from_regular.py

@@ -64,7 +64,7 @@ def action(layout, depth, depth_context, **kwargs):
                 return layout.to_ListOffsetArray64(False)
             elif posaxis == depth and layout.is_list:
                 return layout
-            elif posaxis == 0:
+            elif layout.is_leaf:
                 raise ak._errors.wrap_error(
                     np.AxisError(
                         f"axis={axis} exceeds the depth of this array ({depth})"

src/awkward/operations/ak_is_none.py

@@ -75,7 +75,7 @@ def getfunction_outer(layout, depth, depth_context, **kwargs):
     max_axis = layout.branch_depth[1] - 1
     if axis > max_axis:
         raise ak._errors.wrap_error(
-            np.AxisError(f"axis={axis} exceeds the depth ({max_axis}) of this array")
+            np.AxisError(f"axis={axis} exceeds the depth of this array ({max_axis})")
         )
     behavior = ak._util.behavior_of(array, behavior=behavior)
     depth_context = {"posaxis": axis}