Added ak.firsts / fixed an UnionArray _offsets_and_flattened issue

src/awkward/_v2/contents/unionarray.py

@@ -744,7 +744,7 @@ def _offsets_and_flattened(self, axis, depth):
                 )
 
             else:
-                offsets = ak._v2.index.Index64.zeros(1, self._nplike, dtype=np.int64)
+                offsets = ak._v2.index.Index64.zeros(0, self._nplike, dtype=np.int64)
                 return (
                     offsets,
                     UnionArray(

src/awkward/_v2/operations/structure/ak_firsts.py

@@ -6,62 +6,56 @@
 
 
 def firsts(array, axis=1, highlevel=True, behavior=None):
-    raise ak._v2._util.error(NotImplementedError)
-
-
-#     """
-#     Args:
-#         array: Data from which to select the first elements from nested lists.
-#         axis (int): The dimension at which this operation is applied. The
-#             outermost dimension is `0`, followed by `1`, etc., and negative
-#             values count backward from the innermost: `-1` is the innermost
-#             dimension, `-2` is the next level up, etc.
-#         highlevel (bool): If True, return an #ak.Array; otherwise, return
-#             a low-level #ak.layout.Content subclass.
-#         behavior (None or dict): Custom #ak.behavior for the output array, if
-#             high-level.
-
-#     Selects the first element of each non-empty list and inserts None for each
-#     empty list.
-
-#     For example,
-
-#         >>> array = ak.Array([[1.1], [2.2], [], [3.3], [], [], [4.4], [5.5]])
-#         >>> print(ak.firsts(array))
-#         [1.1, 2.2, None, 3.3, None, None, 4.4, 5.5]
-
-#     See #ak.singletons to invert this function.
-#     """
-#     layout = ak._v2.operations.convert.to_layout(
-#         array, allow_record=False, allow_other=False
-#     )
-#     if isinstance(layout, ak.partition.PartitionedArray):   # NO PARTITIONED ARRAY
-#         posaxis = None
-#         for x in layout.partitions:
-#             if posaxis is None:
-#                 posaxis = x.axis_wrap_if_negative(axis)
-#             elif posaxis != x.axis_wrap_if_negative(axis):
-#                 raise ak._v2._util.error(ValueError(
-#                     "ak.firsts for partitions with different axis depths"
-#
-#                 ))
-#     else:
-#         posaxis = layout.axis_wrap_if_negative(axis)
-
-#     if posaxis == 0:
-#         if len(layout) == 0:
-#             out = None
-#         else:
-#             out = layout[0]
-#     else:
-#         if posaxis < 0:
-#             raise ak._v2._util.error(NotImplementedError(
-#                 "ak.firsts with ambiguous negative axis"
-#
-#             ))
-#         toslice = (slice(None, None, None),) * posaxis + (0,)
-#         out = ak.mask(layout, ak.num(layout, axis=posaxis) > 0, highlevel=False)[
-#             toslice
-#         ]
-
-#     return ak._v2._util.maybe_wrap_like(out, array, behavior, highlevel)
+    """
+    Args:
+        array: Data from which to select the first elements from nested lists.
+        axis (int): The dimension at which this operation is applied. The
+            outermost dimension is `0`, followed by `1`, etc., and negative
+            values count backward from the innermost: `-1` is the innermost
+            dimension, `-2` is the next level up, etc.
+        highlevel (bool): If True, return an #ak.Array; otherwise, return
+            a low-level #ak.layout.Content subclass.
+        behavior (None or dict): Custom #ak.behavior for the output array, if
+            high-level.
+
+    Selects the first element of each non-empty list and inserts None for each
+    empty list.
+
+    For example,
+
+        >>> array = ak.Array([[1.1], [2.2], [], [3.3], [], [], [4.4], [5.5]])
+        >>> print(ak.firsts(array))
+        [1.1, 2.2, None, 3.3, None, None, 4.4, 5.5]
+
+    See #ak.singletons to invert this function.
+    """
+    with ak._v2._util.OperationErrorContext(
+        "ak._v2.firsts",
+        dict(array=array, axis=axis, highlevel=highlevel, behavior=behavior),
+    ):
+        return _impl(array, axis, highlevel, behavior)
+
+
+def _impl(array, axis, highlevel, behavior):
+    layout = ak._v2.operations.convert.to_layout(
+        array, allow_record=False, allow_other=False
+    )
+    posaxis = layout.axis_wrap_if_negative(axis)
+
+    if posaxis == 0:
+        if len(layout) == 0:
+            out = None
+        else:
+            out = layout[0]
+    else:
+        if posaxis < 0:
+            raise ak._v2._util.error(NotImplementedError(
+                "ak.firsts with ambiguous negative axis"
+
+            ))
+        toslice = (slice(None, None, None),) * posaxis + (0,)
+        out = ak._v2.operations.structure.mask(layout, ak._v2.operations.structure.num(layout, axis=posaxis) > 0, highlevel=False)[
+            toslice
+        ]
+
+    return ak._v2._util.wrap(out, behavior, highlevel)

tests/v2/test_0582-propagate-context-in-broadcast_and_apply.py

@@ -120,7 +120,6 @@ def test_cartesian():
         ak._v2.operations.structure.cartesian([one, two], axis=-5)
 
 
-@pytest.mark.skip(reason="ak._v2.operations.structure.firsts to be implemented")
 def test_firsts():
     array = ak._v2.Array([[[0, 1, 2], []], [[3, 4]], [], [[5], [6, 7, 8, 9]]])
 

tests/v2/test_0806-empty-lists-cartesian-fix.py

@@ -19,7 +19,6 @@ def test_empty_arrays_cartesian():
     to_list(ak._v2.operations.structure.concatenate([one, two], axis=0))
 
 
-@pytest.mark.skip(reason="ak._v2.operations.structure.firsts to be implemented")
 def test_cartesian():
     muon = ak._v2.Array([[{"pt": 1.0}], []], with_name="muon")
     electron = ak._v2.Array([[], [{"pt": 1.0}]], with_name="electron")
@@ -42,6 +41,7 @@ def test_cartesian():
     result = ak._v2.operations.structure.cartesian(
         [candidate, ak._v2.Array([[1, 2, 3], []])], axis=1
     )
+
     assert to_list(result) == [None, None]
 
     one, two = ak._v2.operations.structure.broadcast_arrays(