Repository: scikit-hep/awkward

Issue:
`ak.unflatten` error message doesn't indicate axis out of boudns
### Version of Awkward Array

main

### Description and code to reproduce

Given an N-dimensional array for which `ak.unflatten(array, counts, axis=-1)` succeeds, passing `axis>N` does not inform the user that the axis is out of bounds, but rather:
```
Error details: structure imposed by 'counts' does not fit in the array or partition at axis=N
```


Discussion:

Resolving patch:
diff --git a/src/awkward/_util.py b/src/awkward/_util.py
--- a/src/awkward/_util.py
+++ b/src/awkward/_util.py
@@ -900,8 +900,10 @@ def visitor(left, right) -> bool:
             return False
 
         if left.is_list:
-            return numpy.array_equal(left.offsets, right.offsets) and visitor(
-                left.content, right.content
+            return (
+                numpy.array_equal(left.starts, right.starts)
+                and numpy.array_equal(left.stops, right.stops)
+                and visitor(left.content, right.content)
             )
         elif left.is_regular:
             return (left.size == right.size) and visitor(left.content, right.content)
diff --git a/src/awkward/operations/ak_unflatten.py b/src/awkward/operations/ak_unflatten.py
--- a/src/awkward/operations/ak_unflatten.py
+++ b/src/awkward/operations/ak_unflatten.py
@@ -84,21 +84,25 @@ def unflatten(array, counts, axis=0, *, highlevel=True, behavior=None):
 
 
 def _impl(array, counts, axis, highlevel, behavior):
-    layout = ak.operations.to_layout(array, allow_record=False, allow_other=False)
+    layout = ak.operations.to_layout(
+        array, allow_record=False, allow_other=False
+    ).to_packed()
     behavior = ak._util.behavior_of(array, behavior=behavior)
     backend = layout.backend
 
     if isinstance(counts, (numbers.Integral, np.integer)):
         current_offsets = None
     else:
-        counts = ak.operations.to_layout(counts, allow_record=False, allow_other=False)
-
-        if counts.is_option:
+        counts = ak.operations.to_layout(
+            counts, allow_record=False, allow_other=False
+        ).to_packed()
+        if counts.is_option and (counts.content.is_numpy or counts.content.is_unknown):
             mask = counts.mask_as_bool(valid_when=False)
-            counts = counts.to_numpy(allow_missing=True)
-            counts = ak._nplikes.numpy.ma.filled(counts, 0)
+            counts = backend.nplike.to_rectilinear(
+                ak.operations.fill_none(counts, 0, axis=-1, highlevel=False)
+            )
         elif counts.is_numpy or counts.is_unknown:
-            counts = counts.to_numpy(allow_missing=False)
+            counts = backend.nplike.to_rectilinear(counts)
             mask = False
         else:
             raise ak._errors.wrap_error(
@@ -113,11 +117,13 @@ def _impl(array, counts, axis, highlevel, behavior):
         if not issubclass(counts.dtype.type, np.integer):
             raise ak._errors.wrap_error(ValueError("counts must be integers"))
 
-        current_offsets = [backend.index_nplike.empty(len(counts) + 1, np.int64)]
-        current_offsets[0][0] = 0
-        backend.index_nplike.cumsum(counts, out=current_offsets[0][1:])
+        current_offsets = backend.index_nplike.empty(len(counts) + 1, np.int64)
+        current_offsets[0] = 0
+        backend.index_nplike.cumsum(counts, out=current_offsets[1:])
+
+    def unflatten_this_layout(layout):
+        nonlocal current_offsets
 
-    def doit(layout):
         if isinstance(counts, (numbers.Integral, np.integer)):
             if counts < 0 or counts > len(layout):
                 raise ak._errors.wrap_error(
@@ -128,15 +134,15 @@ def doit(layout):
         else:
             position = (
                 backend.index_nplike.searchsorted(
-                    current_offsets[0],
+                    current_offsets,
                     backend.index_nplike.array([len(layout)]),
                     side="right",
                 )[0]
                 - 1
             )
-            if position >= len(current_offsets[0]) or current_offsets[0][
-                position
-            ] != len(layout):
+            if position >= len(current_offsets) or current_offsets[position] != len(
+                layout
+            ):
                 raise ak._errors.wrap_error(
                     ValueError(
                         "structure imposed by 'counts' does not fit in the array or partition "
@@ -144,8 +150,8 @@ def doit(layout):
                     )
                 )
 
-            offsets = current_offsets[0][: position + 1]
-            current_offsets[0] = current_offsets[0][position:] - len(layout)
+            offsets = current_offsets[: position + 1]
+            current_offsets = current_offsets[position:] - len(layout)
 
             out = ak.contents.ListOffsetArray(ak.index.Index64(offsets), layout)
             if not isinstance(mask, (bool, np.bool_)):
@@ -157,23 +163,20 @@ def doit(layout):
         return out
 
     if axis == 0 or ak._util.maybe_posaxis(layout, axis, 1) == 0:
-        out = doit(layout)
+        out = unflatten_this_layout(layout)
 
     else:
 
-        def transform(layout, depth, axis):
-            # Pack the current layout. This ensures that the `counts` array,
-            # which is computed with these layouts applied, aligns with the
-            # internal layout to be unflattened (#910)
-            layout = layout.to_packed()
-
+        def apply(layout, depth, **kwargs):
             posaxis = ak._util.maybe_posaxis(layout, axis, depth)
             if posaxis == depth and layout.is_list:
                 # We are one *above* the level where we want to apply this.
                 listoffsetarray = layout.to_ListOffsetArray64(True)
                 outeroffsets = backend.index_nplike.asarray(listoffsetarray.offsets)
 
-                content = doit(listoffsetarray.content[: outeroffsets[-1]])
+                content = unflatten_this_layout(
+                    listoffsetarray.content[: outeroffsets[-1]]
+                )
                 if isinstance(content, ak.contents.ByteMaskedArray):
                     inneroffsets = backend.index_nplike.asarray(content.content.offsets)
                 elif isinstance(content, ak.contents.RegularArray):
@@ -202,13 +205,10 @@ def transform(layout, depth, axis):
 
                 return ak.contents.ListOffsetArray(ak.index.Index64(positions), content)
 
-            else:
-                return layout
-
-        out = transform(layout, depth=1, axis=axis)
+        out = ak._do.recursively_apply(layout, apply)
 
     if current_offsets is not None and not (
-        len(current_offsets[0]) == 1 and current_offsets[0][0] == 0
+        len(current_offsets) == 1 and current_offsets[0] == 0
     ):
         raise ak._errors.wrap_error(
             ValueError(


Test patch:
diff --git a/tests/test_2071_unflatten_non_packed_counts.py b/tests/test_2071_unflatten_non_packed_counts.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2071_unflatten_non_packed_counts.py
@@ -0,0 +1,35 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest  # noqa: F401
+
+import awkward as ak
+
+
+def test_indexed_counts():
+    counts = ak.contents.IndexedArray(
+        ak.index.Index64(np.arange(3)),
+        ak.contents.NumpyArray(np.array([3, 0, 2], dtype=np.int64)),
+    )
+    assert ak._util.arrays_approx_equal(
+        ak.unflatten([1.1, 2.2, 3.3, 4.4, 5.5], counts),
+        [[1.1, 2.2, 3.3], [], [4.4, 5.5]],
+    )
+
+
+def test_indexed_layout():
+    layout = ak.contents.IndexedArray(
+        ak.index.Index64(np.arange(5)),
+        ak.contents.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=np.float64)),
+    )
+    assert ak._util.arrays_approx_equal(
+        ak.unflatten(layout, [3, 0, 2]),
+        [[1.1, 2.2, 3.3], [], [4.4, 5.5]],
+    )
+
+
+def test_option_counts():
+    assert ak._util.arrays_approx_equal(
+        ak.unflatten([1.1, 2.2, 3.3, 4.4, 5.5], [None, 3, None, 0, 2]),
+        [None, [1.1, 2.2, 3.3], None, [], [4.4, 5.5]],
+    )