1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import numpy as np # noqa: F401
import pytest
import awkward as ak
def test():
original = ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]])
assert ak.operations.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).to_list() == [
[[1, 2], [3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
assert ak.operations.unflatten(original, [1, 3, 1, 2, 1, 1], axis=1).to_list() == [
[[1], [2, 3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
with pytest.raises(ValueError):
ak.operations.unflatten(original, [2, 1, 2, 2, 1, 1], axis=1)
assert ak.operations.unflatten(
original, [2, 0, 2, 1, 2, 1, 1], axis=1
).to_list() == [
[[1, 2], [], [3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
def test_issue742():
assert ak.operations.unflatten(ak.Array(["a", "b", "c"]), [1, 2, 0]).to_list() == [
["a"],
["b", "c"],
[],
]
|