Repository: scikit-hep/awkward

Issue:
unflatten raises counts must be an integer or a one-dimensional array of integers on a 1D awkward-array (worked in 1.x)
### Version of Awkward Array

2.0.5

### Description and code to reproduce

I'm trying to transition from awkward 1.x -> 2.x, and ran into an issue with `ak.unflatten` where the call in 1.x works, but fails in 2.x. In my code, I call:

```python
    num_constituents = ak.num(jets.constituents, axis=2)
    # Convert to offsets
    offsets_constituents = np.cumsum(np.asarray(ak.flatten(num_constituents, axis=1)))
    offsets_constituents = np.insert(offsets_constituents, 0, 0)
    # Then into starts and stops
    starts_constituents = ak.unflatten(offsets_constituents[:-1], ak.num(num_constituents, axis=1))
    stops_constituents = ak.unflatten(offsets_constituents[1:], ak.num(num_constituents, axis=1))
```

For a bit of context, I'm trying to convert some info out from c++ back to python, and have found it convenient to return flattened np arrays and then unflatten on return. In any case, `ak.num(num_constituents, axis=1)` is 1D, as far as I can tell / recall, but in awkward 2.x, I get the following ValueError:

```python
E               ValueError: while calling
E
E                   ak.unflatten(
E                       array = numpy.ndarray([  0   6  37 ... 397 422 442])
E                       counts = <Array [1, 1, 1, 1, 1, ..., 1, 1, 1, 2, 1] type='17 * int64'>
E                       axis = 0
E                       highlevel = True
E                       behavior = None
E                   )
E
E               Error details: counts must be an integer or a one-dimensional array of integers
```

With details on the variables here:

```
array      = array([  0,   6,  37,  50,  81, 118, 146, 173, 193, 208, 242, 275, 318,
       336, 364, 397, 422, 442])
axis       = 0
backend    = <awkward._backends.NumpyBackend object at 0x1109822f0>
behavior   = None
counts     = <IndexedArray len='17'>
    <index><Index dtype='int64' len='17'>
        [ 0  1  3  4  5  6  7  8  9 10 11 12 13 15 16 18 19]
    </Index></index>
    <content><NumpyArray dtype='int64' len='20'>
        [1 1 0 1 1 1 1 1 1 1 1 1 1 1 0 1 1 0 2 1]
    </NumpyArray></content>
</IndexedArray>
highlevel  = True
layout     = <NumpyArray dtype='int64' len='18'>
    [  0   6  37  50  81 118 146 173 193 208 242 275 318 336 364 397 422 442]
</NumpyArray>
```

If I wrap the `ak.num(num_constituents, axis=1)` call with `np.asarray`, my integration tests pass, which at least suggests that it works (although the tests definitely aren't comprehensive), so it may be somewhat trivial. If it just needs to be cast, I think it would be a nicer user experience if `unflatten` did this on it's own. Thanks!



Discussion:
More than a nice user experience, this really should work. All of the array-like function arguments are supposed to be processed with `ak.to_layout`, which would have eliminated any distinction between `np.array(X)` and `ak.Array(np.array(X))`.

I haven't been able to reproduce this. What am I doing differently from you?

```python
>>> import awkward as ak
>>> import numpy as np

>>> ak.unflatten([1.1, 2.2, 3.3, 4.4, 5.5], [3, 0, 2], axis=0)
<Array [[1.1, 2.2, 3.3], [], [4.4, 5.5]] type='3 * var * float64'>

>>> ak.unflatten(ak.Array([1.1, 2.2, 3.3, 4.4, 5.5]), ak.Array([3, 0, 2]), axis=0)
<Array [[1.1, 2.2, 3.3], [], [4.4, 5.5]] type='3 * var * float64'>

>>> ak.unflatten(np.array([1.1, 2.2, 3.3, 4.4, 5.5]), ak.Array([3, 0, 2]), axis=0)
<Array [[1.1, 2.2, 3.3], [], [4.4, 5.5]] type='3 * var * float64'>

>>> ak.unflatten(ak.Array([1.1, 2.2, 3.3, 4.4, 5.5]), np.array([3, 0, 2]), axis=0)
<Array [[1.1, 2.2, 3.3], [], [4.4, 5.5]] type='3 * var * float64'>

>>> ak.unflatten(np.array([1.1, 2.2, 3.3, 4.4, 5.5]), np.array([3, 0, 2]), axis=0)
<Array [[1.1, 2.2, 3.3], [], [4.4, 5.5]] type='3 * var * float64'>
```

Aha! It's the `IndexedArray`:

```python
>>> ak.unflatten([1.1, 2.2, 3.3, 4.4, 5.5], ak.to_categorical([3, 0, 2]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/irishep/awkward/src/awkward/operations/ak_unflatten.py", line 83, in unflatten
    return _impl(array, counts, axis, highlevel, behavior)
  File "/home/jpivarski/irishep/awkward/src/awkward/operations/ak_unflatten.py", line 104, in _impl
    raise ak._errors.wrap_error(
ValueError: while calling

    ak.unflatten(
        array = [1.1, 2.2, 3.3, 4.4, 5.5]
        counts = <Array [3, 0, 2] type='3 * categorical[type=int64]'>
        axis = 0
        highlevel = True
        behavior = None
    )

Error details: counts must be an integer or a one-dimensional array of integers
```

Yes, something is looking at the `counts` and deciding that it's not a flat array of integers because it has an `IndexedArray` node just above the `NumpyArray` node.

https://github.com/scikit-hep/awkward/blob/8d7a2c31f1ea24cef8a909fb12bb14eb812477ac/src/awkward/operations/ak_unflatten.py#L96-L108

The `IndexedArray` case is not included in `counts.is_option`, but it can be handled as a pre-processing. Also, unrelated to the above issue, the `is_option` case should be guarded against failure in `to_numpy` (with a hard-to-interpret error message).

```diff
+        if isinstance(counts, IndexedArray):
+            counts = counts.project()
+
-        if counts.is_option:
+        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)
         elif counts.is_numpy or counts.is_unknown:
             counts = counts.to_numpy(allow_missing=False)
             mask = False
         else:
             raise ak._errors.wrap_error(
                 ValueError(
                     "counts must be an integer or a one-dimensional array of integers"
                 )
             )
```
Thanks for investigating! I'll use `np.asarray` as a workaround, and look out for a future fix (no rush from my side since I haev the workaround). Otherwise, the transition to awkward 2.x looks like it will be pretty easy - I just needed a few minor changes 🙂 
@jpivarski this is closed by #2097, though your categorical test case is not. Originally, I thought that  was out-of-scope for unflatten. I'm perhaps more permissive today. Do you think we should have a special case for categorical integers as `counts`?
I used categorical just as an easy way to produce an IndexedArray, although it's odd that the fix applies to IndexedArrays but not also categoricals. They only differ by a parameter.

I don't see anything in #2097 that explicitly forbids IndexedArrays with `__array__: "categorical"`.
It's more that we project out the index types using `to_packed()` here. If we want to support indexed types, we could. I was thinking of categoricals as a distinct type rather than an optimisation, so it probably does make sense that they work.

Resolving patch:
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
@@ -93,9 +93,10 @@ def _impl(array, counts, axis, highlevel, behavior):
     if isinstance(counts, (numbers.Integral, np.integer)):
         current_offsets = None
     else:
-        counts = ak.operations.to_layout(
-            counts, allow_record=False, allow_other=False
-        ).to_packed()
+        counts = ak.operations.to_layout(counts, allow_record=False, allow_other=False)
+        if counts.is_indexed and not counts.is_option:
+            counts = counts.project()
+
         if counts.is_option and (counts.content.is_numpy or counts.content.is_unknown):
             mask = counts.mask_as_bool(valid_when=False)
             counts = backend.nplike.to_rectilinear(


Test patch:
diff --git a/tests/test_2071_unflatten_non_packed_counts.py b/tests/test_2071_unflatten_non_packed_counts.py
--- a/tests/test_2071_unflatten_non_packed_counts.py
+++ b/tests/test_2071_unflatten_non_packed_counts.py
@@ -33,3 +33,10 @@ def test_option_counts():
         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]],
     )
+
+
+def test_categorical_counts():
+    assert ak._util.arrays_approx_equal(
+        ak.unflatten([1.1, 2.2, 3.3, 4.4, 5.5], ak.to_categorical([3, 0, 2])),
+        [[1.1, 2.2, 3.3], [], [4.4, 5.5]],
+    )