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 47 48 49 50 51 52 53 54 55 56 57 58
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import awkward as ak
def check(layout, assert_length):
if hasattr(layout, "contents"):
for x in layout.contents:
check(x, assert_length)
elif hasattr(layout, "content"):
check(layout.content, assert_length)
else:
assert layout.length <= assert_length
def test_2arrays():
one_a = ak.Array([{"x": 1, "y": 2}], with_name="T")
one_b = ak.Array([{"x": 1, "y": 2}], with_name="T")
two_a = ak.Array([{"x": 1, "z": 3}], with_name="T")
two_b = ak.Array([{"x": 1, "z": 3}], with_name="T")
three = ak.Array([{"x": 4}, {"x": 4}], with_name="T")
first = ak.zip({"a": one_a, "b": one_b})
second = ak.zip({"a": two_a, "b": two_b})
cat = ak.concatenate([first, second], axis=0)
cat["another"] = three
for _ in range(5):
check(cat.layout, 2)
cat["another", "w"] = three.x
def test_3arrays():
zero_a = ak.Array([{"x": 1, "y": 1}], with_name="T")
zero_b = ak.Array([{"x": 1, "v": 1}], with_name="T")
one_a = ak.Array([{"x": 1, "y": 2}], with_name="T")
one_b = ak.Array([{"x": 1, "y": 2}], with_name="T")
two_a = ak.Array([{"x": 1, "z": 3}], with_name="T")
two_b = ak.Array([{"x": 1, "z": 3}], with_name="T")
three = ak.Array([{"x": 4}, {"x": 4}, {"x": 4}], with_name="T")
zeroth = ak.zip({"a": zero_a, "b": zero_b})
first = ak.zip({"a": one_a, "b": one_b})
second = ak.zip({"a": two_a, "b": two_b})
cat = ak.concatenate([zeroth, first, second], axis=0)
cat["another"] = three
for _ in range(5):
check(cat.layout, 3)
cat["another", "w"] = three.x
|