fix: handle type tracer correctly in indexed option array flattening (#…

…3325) * fix: handle type tracer correctly in indexed option array * add test * lint * lint * lint --------- Co-authored-by: Jim Pivarski <[email protected]>
scikit-hep · Dec 5, 2024 · 7191ce8 · 7191ce8
1 parent 7d97988
commit 7191ce8
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 1 deletion.
diff --git a/src/awkward/contents/indexedoptionarray.py b/src/awkward/contents/indexedoptionarray.py
@@ -600,7 +600,7 @@ def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content]
 
             offsets, flattened = next._offsets_and_flattened(axis, depth)
 
-            if offsets.length == 0:
+            if offsets.length is not unknown_length and offsets.length == 0:
                 return (
                     offsets,
                     ak.contents.IndexedOptionArray(

diff --git a/tests/test_3325_ak_flatten_indexed_option_array.py b/tests/test_3325_ak_flatten_indexed_option_array.py
@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import pytest
+
+import awkward as ak
+
+fromdict = {
+    "class": "ListOffsetArray",
+    "offsets": "i64",
+    "content": {
+        "class": "ListOffsetArray",
+        "offsets": "i64",
+        "content": {
+            "class": "IndexedOptionArray",
+            "index": "i64",
+            "content": {
+                "class": "ListOffsetArray",
+                "offsets": "i64",
+                "content": {
+                    "class": "IndexedOptionArray",
+                    "index": "i64",
+                    "content": {
+                        "class": "NumpyArray",
+                        "primitive": "float32",
+                        "inner_shape": [],
+                        "parameters": {},
+                        "form_key": None,
+                    },
+                    "parameters": {},
+                    "form_key": None,
+                },
+                "parameters": {},
+                "form_key": None,
+            },
+            "parameters": {},
+            "form_key": None,
+        },
+        "parameters": {},
+        "form_key": None,
+    },
+    "parameters": {},
+    "form_key": None,
+}
+
+form = ak.forms.from_dict(fromdict)
+
+ttlayout, report = ak.typetracer.typetracer_with_report(form)
+
+ttarray = ak.Array(ttlayout)
+
+
+@pytest.mark.parametrize("ax", [None, 0, 1, 2, 3])
+def test_3325_flatten_index_option_array(ax):
+    ak.flatten(ttarray, axis=ax)