Repository: scikit-hep/awkward

Issue:
IndexError: boolean index did not match indexed array in awkward 2.x vs 1.x
### Version of Awkward Array

2.0.6

### Description and code to reproduce

Following up on my previous report, I've found another case where the behavior varies from awkward 1.x to 2.x in an unexpected manner. The rough outline is that I'm trying to apply a mask to select some subset of nested data. In awkward 1.x, this works fine (I did this with a behavior, which you can see a glimpse of the functionality in the trace - here it's a bit for convenience). In awkward 2.x, this raises an IndexError (trace below):

```python
Cell In [6], line 1
----> 1 jets.jet_splittings.iterative_splittings(jets.subjets)

File ~/software/dev/mammoth/src/mammoth/framework/analysis/jet_substructure.py:387, in JetSplittingArray.iterative_splittings(self, subjets)
    379 def iterative_splittings(self, subjets: SubjetArray) -> SubjetArray:
    380     """Retrieve the iterative splittings.
    381
    382     Args:
   (...)
    385         The splittings which are part of the iterative splitting chain.
    386     """
--> 387     return cast(SubjetArray, self[subjets.iterative_splitting_index])

File ~/software/dev/mammoth/src/mammoth/framework/analysis/jet_substructure.py:253, in SubjetArray.iterative_splitting_index(self)
    250 @property
    251 def iterative_splitting_index(self) -> AwkwardArray[int]:
    252     """Indices of splittings which were part of the iterative splitting chain."""
--> 253     return self.parent_splitting_index[self.part_of_iterative_splitting]

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/highlevel.py:956, in Array.__getitem__(self, where)
    527 """
    528 Args:
    529     where (many types supported; see below): Index of positions to
   (...)
    953 have the same dimension as the array being indexed.
    954 """
    955 with ak._errors.SlicingErrorContext(self, where):
--> 956     out = self._layout[where]
    957     if isinstance(out, ak.contents.NumpyArray):
    958         array_param = out.parameter("__array__")

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/contents/content.py:509, in Content.__getitem__(self, where)
    508 def __getitem__(self, where):
--> 509     return self._getitem(where)

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/contents/content.py:554, in Content._getitem(self, where)
    551         return out._getitem_at(0)
    553 elif isinstance(where, ak.highlevel.Array):
--> 554     return self._getitem(where.layout)
    556 elif (
    557     isinstance(where, Content)
    558     and where._parameters is not None
    559     and (where._parameters.get("__array__") in ("string", "bytestring"))
    560 ):
    561     return self._getitem_fields(ak.operations.to_list(where))

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/contents/content.py:602, in Content._getitem(self, where)
    599         return out._getitem_at(0)
    601 elif isinstance(where, Content):
--> 602     return self._getitem((where,))
    604 elif ak._util.is_sized_iterable(where) and len(where) == 0:
    605     return self._carry(
    606         ak.index.Index64.empty(0, self._backend.index_nplike),
    607         allow_lazy=True,
    608     )

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/contents/content.py:535, in Content._getitem(self, where)
    532     return self
    534 # Normalise valid indices onto well-defined basis
--> 535 items = ak._slicing.normalise_items(where, self._backend)
    536 # Prepare items for advanced indexing (e.g. via broadcasting)
    537 nextwhere = ak._slicing.prepare_advanced_indexing(items)

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/_slicing.py:206, in normalise_items(where, backend)
    204 def normalise_items(where: Sequence, backend: ak._backends.Backend) -> list:
    205     # First prepare items for broadcasting into like-types
--> 206     return [normalise_item(x, backend) for x in where]

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/_slicing.py:206, in <listcomp>(.0)
    204 def normalise_items(where: Sequence, backend: ak._backends.Backend) -> list:
    205     # First prepare items for broadcasting into like-types
--> 206     return [normalise_item(x, backend) for x in where]

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/_slicing.py:172, in normalise_item(item, backend)
    169         return as_numpy.data
    171 elif isinstance(item, ak.contents.Content):
--> 172     out = normalise_item_bool_to_int(normalise_item_nested(item))
    173     if isinstance(out, ak.contents.NumpyArray):
    174         return out.data

File ~/software/dev/mammoth/.venv/lib/python3.10/site-packages/awkward/_slicing.py:373, in normalise_item_bool_to_int(item)
    371 if item.backend.nplike.known_data or item.backend.nplike.known_shape:
    372     localindex = ak._do.local_index(item, axis=1)
--> 373     nextcontent = localindex.content.data[item.content.data]
    375     cumsum = item.backend.index_nplike.empty(
    376         item.content.data.shape[0] + 1, np.int64
    377     )
    378     cumsum[0] = 0

IndexError: boolean index did not match indexed array along dimension 0; dimension is 14 but corresponding boolean dimension is 36
```

Unfortunately, this has been tricky to reproduce since it's deeply nested in my code, and only seems to occur for some inputs (there's also a random input :-( ). I tried pickling the array, but in testing the pickled array, it works fine. This led me to guessing a workaround - that if I `ak.to_packed()` the data, it would work, and anecdotally, this seemed to have worked (in some previous issue, I learned that pickling calls something like `to_packed`). I managed to reproduce it once in a debug environment (unfortunately, unclear if it will work consistently - there are still some random elements), and grabbed the layouts from the initial array (failing) and on the packed array (working)

original:
```
In [15]: jets.layout
Out[15]:
<RecordArray is_tuple='false' len='1'>
    <content index='0' field='jet_pt'>
        <NumpyArray dtype='float64' len='1'>[24.54459219]</NumpyArray>
    </content>
    <content index='1' field='jet_constituents'>
        <ListOffsetArray len='1'>
            <offsets><Index dtype='int64' len='2'>
                [0 8]
            </Index></offsets>
            <content><RecordArray is_tuple='false' len='8'>
                <parameter name='__record__'>'JetConstituent'</parameter>
                <content index='0' field='pt'>
                    <NumpyArray dtype='float64' len='8'>
                        [ 1.54625939  0.76237501  0.23116202  1.25449706
                          2.66868935  5.11290739  0.70627872 12.29195148]
                    </NumpyArray>
                </content>
                <content index='1' field='eta'>
                    <NumpyArray dtype='float64' len='8'>
                        [0.53683598 0.82056333 0.84215039 0.76359227 0.59864123
                         0.69776704 0.72732497 0.69191189]
                    </NumpyArray>
                </content>
                <content index='2' field='phi'>
                    <NumpyArray dtype='float64' len='8'>
                        [-2.43178638 -2.41334926 -2.50400814 -2.41215308
                         -2.58415947 -2.55853468 -2.52816211 -2.49251164]
                    </NumpyArray>
                </content>
                <content index='3' field='id'>
                    <IndexedArray len='8'>
                        <index><Index dtype='int64' len='8'>
                            [726608 726112 726550 726268 726302 726309 726566
                             726301]
                        </Index></index>
                        <content><NumpyArray dtype='int64' len='1420975'>
                            [     0 100000 100001 ... 100664 100665 100666]
                        </NumpyArray></content>
                    </IndexedArray>
                </content>
            </RecordArray></content>
        </ListOffsetArray>
    </content>
    <content index='2' field='jet_splittings'>
        <IndexedArray len='1'>
            <index><Index dtype='int64' len='1'>
                [0]
            </Index></index>
            <content><ListOffsetArray len='1'>
                <offsets><Index dtype='int64' len='2'>
                    [0 7]
                </Index></offsets>
                <content><RecordArray is_tuple='false' len='7'>
                    <parameter name='__record__'>'JetSplitting'</parameter>
                    <content index='0' field='kt'>
                        <NumpyArray dtype='float64' len='7'>
                            [0.27029765 0.32986462 0.31952325 0.36662835
                             0.02992396 0.02485513 0.04341938]
                        </NumpyArray>
                    </content>
                    <content index='1' field='delta_R'>
                        <NumpyArray dtype='float64' len='7'>
                            [0.17571019 0.14732361 0.12001833 0.06304824
                             0.04238117 0.10773081 0.05698362]
                        </NumpyArray>
                    </content>
                    <content index='2' field='z'>
                        <NumpyArray dtype='float64' len='7'>
                            [0.06298601 0.09764745 0.12847613 0.32129356
                             0.1213707  0.10282853 0.37799871]
                        </NumpyArray>
                    </content>
                    <content index='3' field='parent_index'>
                        <NumpyArray dtype='int64' len='7'>[-1  0  1  2  3  1  5]</NumpyArray>
                    </content>
                </RecordArray></content>
            </ListOffsetArray></content>
        </IndexedArray>
    </content>
    <content index='3' field='subjets'>
        <ListOffsetArray len='1'>
            <offsets><Index dtype='int64' len='2'>
                [22 36]
            </Index></offsets>
            <content><RecordArray is_tuple='false' len='36'>
                <parameter name='__record__'>'Subjet'</parameter>
                <content index='0' field='part_of_iterative_splitting'>
                    <NumpyArray dtype='bool' len='36'>
                        [ True False  True False  True False  True False  True
                         False  True False False False False False False False
                         False False False False  True False  True False  True
                         False  True False False False False False False False]
                    </NumpyArray>
                </content>
                <content index='1' field='parent_splitting_index'>
                    <NumpyArray dtype='int64' len='36'>
                        [ 0  0  1  1  2  2  3  3  4  4  5  5  6  6  7  7  8  8
                          9  9 10 10  0  0  1  1  2  2  3  3  4  4  5  5  6  6]
                    </NumpyArray>
                </content>
                <content index='2' field='constituent_indices'>
                    <ListOffsetArray len='36'>
                        <offsets><Index dtype='int64' len='37'>
                            [ 0 11 12 21 23 29 32 35 38 40 41 42 43 44 46 47 48
                             50 51 52 53 54 55 62 63 67 70 73 74 75 77 78 79 81
                             82 83 84]
                        </Index></offsets>
                        <content><NumpyArray dtype='int64' len='84'>
                            [ 1  3  7  2  6  8  4  5  9 10 11  0  7  2  6  8  4
                              5  9 10 11  1  3  8  4  5  9 10 11  7  2  6  9 10
                             11  8  4  5 10 11  9 10 11  8  4  5  5  4  2  6  7
                              2  6  3  1  2  1  3  4  7  5  6  0  4  7  5  6  2
                              1  3  7  5  6  4  7  5  6  5  6  1  3  2  3  1]
                        </NumpyArray></content>
                    </ListOffsetArray>
                </content>
            </RecordArray></content>
        </ListOffsetArray>
    </content>
    <content index='4' field='leading_track_pt'>
        <IndexedArray len='1'>
            <index><Index dtype='int64' len='1'>
                [1]
            </Index></index>
            <content><NumpyArray dtype='float64' len='2'>[12.69099712 12.34701824]</NumpyArray></content>
        </IndexedArray>
    </content>
</RecordArray>
```

packed:

```
In [16]: jets_packed = ak.to_packed(jets)

In [17]: jets_packed.layout
Out[17]:
<RecordArray is_tuple='false' len='1'>
    <content index='0' field='jet_pt'>
        <NumpyArray dtype='float64' len='1'>[24.54459219]</NumpyArray>
    </content>
    <content index='1' field='jet_constituents'>
        <ListOffsetArray len='1'>
            <offsets><Index dtype='int64' len='2'>
                [0 8]
            </Index></offsets>
            <content><RecordArray is_tuple='false' len='8'>
                <parameter name='__record__'>'JetConstituent'</parameter>
                <content index='0' field='pt'>
                    <NumpyArray dtype='float64' len='8'>
                        [ 1.54625939  0.76237501  0.23116202  1.25449706
                          2.66868935  5.11290739  0.70627872 12.29195148]
                    </NumpyArray>
                </content>
                <content index='1' field='eta'>
                    <NumpyArray dtype='float64' len='8'>
                        [0.53683598 0.82056333 0.84215039 0.76359227 0.59864123
                         0.69776704 0.72732497 0.69191189]
                    </NumpyArray>
                </content>
                <content index='2' field='phi'>
                    <NumpyArray dtype='float64' len='8'>
                        [-2.43178638 -2.41334926 -2.50400814 -2.41215308
                         -2.58415947 -2.55853468 -2.52816211 -2.49251164]
                    </NumpyArray>
                </content>
                <content index='3' field='id'>
                    <NumpyArray dtype='int64' len='8'>
                        [100494      5 100436 100154 100188 100195 100452
                         100187]
                    </NumpyArray>
                </content>
            </RecordArray></content>
        </ListOffsetArray>
    </content>
    <content index='2' field='jet_splittings'>
        <ListOffsetArray len='1'>
            <offsets><Index dtype='int64' len='2'>[0 7]</Index></offsets>
            <content><RecordArray is_tuple='false' len='7'>
                <parameter name='__record__'>'JetSplitting'</parameter>
                <content index='0' field='kt'>
                    <NumpyArray dtype='float64' len='7'>
                        [0.27029765 0.32986462 0.31952325 0.36662835 0.02992396
                         0.02485513 0.04341938]
                    </NumpyArray>
                </content>
                <content index='1' field='delta_R'>
                    <NumpyArray dtype='float64' len='7'>
                        [0.17571019 0.14732361 0.12001833 0.06304824 0.04238117
                         0.10773081 0.05698362]
                    </NumpyArray>
                </content>
                <content index='2' field='z'>
                    <NumpyArray dtype='float64' len='7'>
                        [0.06298601 0.09764745 0.12847613 0.32129356 0.1213707
                         0.10282853 0.37799871]
                    </NumpyArray>
                </content>
                <content index='3' field='parent_index'>
                    <NumpyArray dtype='int64' len='7'>[-1  0  1  2  3  1  5]</NumpyArray>
                </content>
            </RecordArray></content>
        </ListOffsetArray>
    </content>
    <content index='3' field='subjets'>
        <ListOffsetArray len='1'>
            <offsets><Index dtype='int64' len='2'>[ 0 14]</Index></offsets>
            <content><RecordArray is_tuple='false' len='14'>
                <parameter name='__record__'>'Subjet'</parameter>
                <content index='0' field='part_of_iterative_splitting'>
                    <NumpyArray dtype='bool' len='14'>
                        [ True False  True False  True False  True False False
                         False False False False False]
                    </NumpyArray>
                </content>
                <content index='1' field='parent_splitting_index'>
                    <NumpyArray dtype='int64' len='14'>[0 0 1 1 2 2 3 3 4 4 5 5 6 6]</NumpyArray>
                </content>
                <content index='2' field='constituent_indices'>
                    <ListOffsetArray len='14'>
                        <offsets><Index dtype='int64' len='15'>
                            [ 0  7  8 12 15 18 19 20 22 23 24 26 27 28 29]
                        </Index></offsets>
                        <content><NumpyArray dtype='int64' len='29'>
                            [2 1 3 4 7 5 6 0 4 7 5 6 2 1 3 7 5 6 4 7 5 6 5 6 1
                             3 2 3 1]
                        </NumpyArray></content>
                    </ListOffsetArray>
                </content>
            </RecordArray></content>
        </ListOffsetArray>
    </content>
    <content index='4' field='leading_track_pt'>
        <NumpyArray dtype='float64' len='1'>[12.34701824]</NumpyArray>
    </content>
</RecordArray>
```

`to_packed` seems like it's given the right behavior since the original and packed jets appear to have the same values, but I think it's a bug that it's needed here. Can you please take a look?

(I'm not sure what the performance implications of `ak.to_packed` are here, but I need a quick fix, and it seems to work, so I'll run with it. If there are cheaper/better workarounds, I would appreciate any suggestions!)


Discussion:

Resolving patch:
diff --git a/src/awkward/_slicing.py b/src/awkward/_slicing.py
--- a/src/awkward/_slicing.py
+++ b/src/awkward/_slicing.py
@@ -440,6 +440,7 @@ def _normalise_item_bool_to_int(item: Content, backend: Backend) -> Content:
         and np.issubdtype(item.content.dtype, np.bool_)
     ):
         if item_backend.nplike.known_data or item_backend.nplike.known_shape:
+            item = item.to_ListOffsetArray64(True)
             localindex = ak._do.local_index(item, axis=1)
             nextcontent = localindex.content.data[item.content.data]
 


Test patch:
diff --git a/tests/test_2214_offset_bool_index.py b/tests/test_2214_offset_bool_index.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2214_offset_bool_index.py
@@ -0,0 +1,38 @@
+# 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():
+    # fmt: off
+    layout = ak.contents.RecordArray(
+        [
+            ak.contents.ListOffsetArray(
+                ak.index.Index64(np.array([22, 36], dtype=np.int64)),
+                ak.contents.RecordArray(
+                    [
+                        ak.contents.NumpyArray(
+                            np.array([True, False, True, False, True, False, True, False, True, False, True, False, False, False,
+                             False, False, False, False, False, False, False, False, True, False, True, False, True,
+                             False, True, False, False, False, False, False, False, False], dtype=np.bool_)
+                        ),
+                        ak.contents.NumpyArray(
+                            np.array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6], dtype=np.int64)
+                        )
+                    ],
+                    ["part_of_iterative_splitting", "parent_splitting_index"],
+                ),
+            )
+        ],
+        ["subjets"],
+    )
+    # fmt: on
+
+    array = ak.Array(layout)
+    result = array.subjets.parent_splitting_index[
+        array.subjets.part_of_iterative_splitting
+    ]
+    assert ak.almost_equal(result, [[0, 1, 2, 3]])