From 7191ce8247c810eb5ebffc0b31a076583939f22d Mon Sep 17 00:00:00 2001
From: Lindsey Gray <lindsey.gray@gmail.com>
Date: Thu, 5 Dec 2024 14:23:34 -0600
Subject: [PATCH] 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 <jpivarski@users.noreply.github.com>
---
 src/awkward/contents/indexedoptionarray.py    |  2 +-
 ...st_3325_ak_flatten_indexed_option_array.py | 54 +++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_3325_ak_flatten_indexed_option_array.py

diff --git a/src/awkward/contents/indexedoptionarray.py b/src/awkward/contents/indexedoptionarray.py
index 57e67c6aa9..dc67d1ca3f 100644
--- 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
new file mode 100644
index 0000000000..a8697a6993
--- /dev/null
+++ 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)