fix: ak.flatten on BitMaskedArray (#3034)

scikit-hep · Feb 26, 2024 · 6a9c691 · 6a9c691
1 parent bcc9509
commit 6a9c691
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/src/awkward/contents/bitmaskedarray.py b/src/awkward/contents/bitmaskedarray.py
@@ -570,7 +570,7 @@ def project(self, mask=None):
         return self.to_ByteMaskedArray().project(mask)
 
     def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content]:
-        return self.to_ByteMaskedArray._offsets_and_flattened(axis, depth)
+        return self.to_ByteMaskedArray()._offsets_and_flattened(axis, depth)
 
     def _mergeable_next(self, other: Content, mergebool: bool) -> bool:
         # Is the other content is an identity, or a union?

diff --git a/tests/test_3033_flatten_bitmaskedarray.py b/tests/test_3033_flatten_bitmaskedarray.py
@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import numpy as np
+
+import awkward as ak
+
+
+def test():
+    layout = ak.to_layout([[[1, 2, 3], [4]], None]).to_BitMaskedArray(
+        valid_when=True, lsb_order=True
+    )
+    result = ak.flatten(layout, axis=1, highlevel=False)
+    assert result.is_equal_to(
+        ak.contents.ListOffsetArray(
+            ak.index.Index64([0, 3, 4]),
+            ak.contents.NumpyArray(np.array([1, 2, 3, 4], dtype=np.int64)),
+        )
+    )