fix: make to_packed work for typetracer backed IndexedOptionArray, BitMaskedArray, and UnionArray (#3608)

* make to_packed work for indexedoptionarray with typetracer

* add test

* better use known_data

* fix other layouts too and add more tests

* better like this

* guard the legnth, not the nplike

---------

Co-authored-by: Ianna Osborne <ianna.osborne@cern.ch>

Author: ikrommyd

src/awkward/contents/bitmaskedarray.py

@@ -810,9 +810,13 @@ def _to_packed(self, recursive: bool = True) -> Self:
             )
 
         else:
-            excess_length = math.ceil(self.length / 8.0)
+            if self.length is not unknown_length:
+                excess_length = math.ceil(self.length / 8.0)
+            else:
+                excess_length = unknown_length
             if (
                 self._mask.length is not unknown_length
+                and excess_length is not unknown_length
                 and self._mask.length == excess_length
             ):
                 mask = self._mask

src/awkward/contents/bytemaskedarray.py

@@ -335,7 +335,7 @@ def to_ByteMaskedArray(self, valid_when):
     def to_BitMaskedArray(self, valid_when, lsb_order):
         if not self._backend.nplike.known_data:
             self._touch_data(recursive=False)
-            if self._backend.nplike.known_data:
+            if self.length is not unknown_length:
                 excess_length = math.ceil(self.length / 8.0)
             else:
                 excess_length = unknown_length

src/awkward/contents/indexedoptionarray.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import copy
-import operator
 from collections.abc import Mapping, MutableMapping, Sequence
 
 import awkward as ak
@@ -1722,7 +1721,7 @@ def _to_packed(self, recursive: bool = True) -> Self:
             nplike = self._backend.nplike
             original_index = self._index.data
             is_none = original_index < 0
-            num_none = operator.index(nplike.count_nonzero(is_none))
+            num_none = nplike.index_as_shape_item(nplike.count_nonzero(is_none))
             new_index = nplike.empty(self._index.length, dtype=self._index.dtype)
             if isinstance(nplike, Jax):
                 new_index = new_index.at[is_none].set(-1)

src/awkward/contents/unionarray.py

@@ -1674,15 +1674,21 @@ def continuation():
     def _to_packed(self, recursive: bool = True) -> Self:
         nplike = self._backend.nplike
         tags = self._tags.data
-        original_index = index = self._index.data[: tags.shape[0]]
+        original_index = index = self._index.data[
+            : nplike.shape_item_as_index(tags.shape[0])
+        ]
 
         contents = list(self._contents)
 
         for tag in range(len(self._contents)):
             is_tag = tags == tag
             num_tag = nplike.index_as_shape_item(nplike.count_nonzero(is_tag))
 
-            if len(contents[tag]) > num_tag:
+            if (
+                contents[tag].length is not unknown_length
+                and num_tag is not unknown_length
+                and contents[tag].length > num_tag
+            ):
                 if original_index is index:
                     index = index.copy()
                 new_index_values = self._backend.nplike.arange(

tests/test_3608_to_packed_for_typetracer_backed_arrays.py

@@ -0,0 +1,163 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import awkward as ak
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_numpy_array(forget_length, recursive):
+    matrix = np.arange(64).reshape(8, -1)
+    layout = ak.contents.NumpyArray(matrix[:, 0])
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_empty_array(forget_length, recursive):
+    layout = ak.contents.EmptyArray()
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_indexed_option_array(forget_length, recursive):
+    index = ak.index.Index64(np.r_[0, -1, 2, -1, 4])
+    content = ak.contents.NumpyArray(np.arange(8))
+    layout = ak.contents.IndexedOptionArray(index, content)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_indexed_array(forget_length, recursive):
+    index = ak.index.Index64(np.array([0, 1, 2, 3, 6, 7, 8]))
+    content = ak.contents.NumpyArray(np.arange(10))
+    layout = ak.contents.IndexedArray(index, content)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_list_array(forget_length, recursive):
+    content = ak.contents.NumpyArray(
+        np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
+    )
+    starts = ak.index.Index64(np.array([0, 3, 3, 5, 6]))
+    stops = ak.index.Index64(np.array([3, 3, 5, 6, 9]))
+    layout = ak.contents.ListArray(starts, stops, content)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_list_offset_array(forget_length, recursive):
+    content = ak.contents.NumpyArray(
+        np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
+    )
+    offsets = ak.index.Index64(np.array([0, 3, 3, 5, 6]))
+    layout = ak.contents.ListOffsetArray(offsets, content)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_unmasked_array(forget_length, recursive):
+    content = ak.contents.NumpyArray(
+        np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
+    )
+    layout = ak.contents.UnmaskedArray(content)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_union_array(forget_length, recursive):
+    a = ak.contents.NumpyArray(np.arange(4))
+    b = ak.contents.NumpyArray(np.arange(4) + 4)
+    c = ak.contents.RegularArray(ak.contents.NumpyArray(np.arange(12)), 3)
+    layout = ak.contents.UnionArray.simplified(
+        ak.index.Index8([1, 1, 2, 2, 0, 0]),
+        ak.index.Index64([0, 1, 0, 1, 0, 1]),
+        [a, b, c],
+    )
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_record_array(forget_length, recursive):
+    a = ak.contents.NumpyArray(np.arange(10))
+    b = ak.contents.NumpyArray(np.arange(10) * 2 + 4)
+    layout = ak.contents.RecordArray([a, b], None, 5)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_regular_array(forget_length, recursive):
+    content = ak.contents.NumpyArray(np.arange(10))
+    layout = ak.contents.RegularArray(content, 3)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_bit_masked_array(forget_length, recursive):
+    mask = ak.index.IndexU8(np.array([0b10101010]))
+    content = ak.contents.NumpyArray(np.arange(16))
+    layout = ak.contents.BitMaskedArray(mask, content, False, 8, False)
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )
+
+
+@pytest.mark.parametrize("forget_length", [True, False])
+@pytest.mark.parametrize("recursive", [True, False])
+def test_byte_masked_array(forget_length, recursive):
+    mask = ak.index.Index8(np.array([1, 0, 1, 0, 1, 0, 1, 0]))
+    content = ak.contents.NumpyArray(np.arange(16))
+    layout = ak.contents.ByteMaskedArray(
+        mask,
+        content,
+        False,
+    )
+    assert (
+        layout.to_typetracer(forget_length).to_packed(recursive).form
+        == layout.to_packed(recursive).form
+    )