Skip to content

Commit

Permalink
fix: support typetracer in to_packed for IndexedOptionArray (#2912)
Browse files Browse the repository at this point in the history
* fix: support typetracer in `to_packed` for `IndexedOptionArray`

* fix: don't avoid touching if typetracer
  • Loading branch information
agoose77 authored Jan 2, 2024
1 parent 0a8e9c8 commit d466d70
Showing 1 changed file with 11 additions and 12 deletions.
23 changes: 11 additions & 12 deletions src/awkward/contents/indexedoptionarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1718,27 +1718,26 @@ def continuation():
raise AssertionError(result)

def to_packed(self) -> Self:
original_index = self._index.raw(self._backend.nplike)
is_none = original_index < 0
num_none = self._backend.index_nplike.count_nonzero(is_none)
if self.parameter("__array__") == "categorical" or self._content.length <= (
len(original_index) - num_none
index_nplike = self._backend.index_nplike
is_none = self._index.data < 0
num_none = index_nplike.index_as_shape_item(index_nplike.count_nonzero(is_none))
if self.parameter("__array__") == "categorical" or (
self._backend.index_nplike.known_data
and self._content.length <= (self._index.length - num_none)
):
return ak.contents.IndexedOptionArray(
self._index, self._content.to_packed(), parameters=self._parameters
)

else:
new_index = self._backend.index_nplike.empty(
len(original_index), dtype=original_index.dtype
)
new_index = index_nplike.empty(self._index.length, dtype=self._index.dtype)
new_index[is_none] = -1
new_index[~is_none] = self._backend.index_nplike.arange(
len(original_index) - num_none,
dtype=original_index.dtype,
new_index[~is_none] = index_nplike.arange(
index_nplike.shape_item_as_index(new_index.size - num_none),
dtype=self._index.dtype,
)
return ak.contents.IndexedOptionArray(
ak.index.Index(new_index, nplike=self._backend.index_nplike),
ak.index.Index(new_index, nplike=index_nplike),
self.project().to_packed(),
parameters=self._parameters,
)
Expand Down

0 comments on commit d466d70

Please sign in to comment.