Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: structure removal order for unions #3397

Merged
merged 14 commits into from
Mar 5, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 15 additions & 5 deletions src/awkward/contents/unionarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1574,13 +1574,23 @@ def _remove_structure(
self, backend: Backend, options: RemoveStructureOptions
) -> list[Content]:
out = []
for i in range(len(self._contents)):
index = self._index[self._tags.data == i]
out.extend(
self._contents[i]
._carry(index, False)

# typetracer
if not self._backend.nplike.known_data:
self._touch_data(recursive=False)
# just flatten, ignore order, tags, index
for c in self._contents:
out.extend(c._remove_structure(backend, options))
return out

# backends with concrete data
for i in range(self._tags.length):
content = (
self._contents[self._tags.data[i]]
._carry(ak.index.Index(self._index.data[i]), False)
._remove_structure(backend, options)
)
out.extend(content)
Comment on lines +1578 to +1593
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@pfackeldey this is a correct solution as far as structure goes, although allocates repeatedly according to the tags which is not ideal.

Would you be able to leave a # FIXME: note here (probably just commit to main) that we can optimise this by building a UnionArray over the original naive loop to re-order the result? That might not be trivial to do, but we should definitely remark upon it!

return out

def _recursively_apply(
Expand Down
62 changes: 62 additions & 0 deletions tests/test_3180_unionarrays_order_structure_removal.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
# ruff: noqa: E402

from __future__ import annotations

import pytest

import awkward as ak


def test_ravel_boolean():
x = ak.Array([None, [False], [True]])
y = ak.ravel(ak.fill_none(x, False, axis=0))
assert y.tolist() == [False, False, True]


def test_ravel_boolean_typetracer():
x = ak.Array([None, [False], [True]], backend="typetracer")
y = ak.ravel(ak.fill_none(x, False, axis=0))
assert y.typestr == "## * bool"


def test_ravel_ListOffsetArray():
a = ak.Array([[1, 2, 3], 4, [5, 6], 7, 8])
b = ak.ravel(a)
assert b.tolist() == [1, 2, 3, 4, 5, 6, 7, 8]


def test_ravel_ListOffsetArray_typetracer():
a = ak.Array([[1, 2, 3], 4, [5, 6], 7, 8], backend="typetracer")
b = ak.ravel(a)
assert b.typestr == "## * int64"


def test_ravel_nested_ListOffsetArrays():
a = ak.Array([[[1, 2], 3], 4, [5, [[[6]]]], 7, 8])
b = ak.ravel(a)
assert b.tolist() == [1, 2, 3, 4, 5, 6, 7, 8]


def test_ravel_nested_ListOffsetArrays_typetracer():
a = ak.Array([[[1, 2], 3], 4, [5, [[[6]]]], 7, 8], backend="typetracer")
b = ak.ravel(a)
assert b.typestr == "## * int64"


def test_ravel_incompatible_contents():
a = ak.Array([1, "high", 2, "low"])
with pytest.raises(
AssertionError,
match="cannot merge NumpyArray with ListOffsetArray",
):
ak.ravel(a)


def test_ravel_incompatible_contents_typetracer():
a = ak.Array([1, "high", 2, "low"], backend="typetracer")
with pytest.raises(
AssertionError,
match="cannot merge NumpyArray with ListOffsetArray",
):
ak.ravel(a)
Loading