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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
# 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
pa = pytest.importorskip("pyarrow")
def test():
one = ak.zip({"pt": [[1, 2], [1]], "eta": [[1.1, 2.1], [1]]})
two = ak.zip({"pt": [[2.1], [2.21, 2.22]], "eta": [[2.11], [2.221, 2.222]]})
assert ak.cartesian({"foo": one, "bar": two}, nested=False).layout.is_equal_to(
ak.contents.ListOffsetArray(
ak.index.Index64([0, 2, 4]),
ak.contents.RecordArray(
[
ak.contents.IndexedArray(
ak.index.Index64([0, 1, 2, 2]),
ak.contents.RecordArray(
[
ak.contents.NumpyArray(
np.array([1, 2, 1], dtype=np.int64)
),
ak.contents.NumpyArray(
np.array([1.1, 2.1, 1.0], dtype=np.float64)
),
],
["pt", "eta"],
),
),
ak.contents.IndexedArray(
ak.index.Index64([0, 0, 1, 2]),
ak.contents.RecordArray(
[
ak.contents.NumpyArray(
np.array([2.1, 2.21, 2.22], dtype=np.float64)
),
ak.contents.NumpyArray(
np.array([2.11, 2.221, 2.222], dtype=np.float64)
),
],
["pt", "eta"],
),
),
],
["foo", "bar"],
),
)
)
assert ak.cartesian({"foo": one, "bar": two}, nested=True).layout.is_equal_to(
ak.contents.ListOffsetArray(
ak.index.Index64([0, 2, 3]),
ak.contents.ListOffsetArray(
ak.index.Index64([0, 1, 2, 4]),
ak.contents.RecordArray(
[
ak.contents.IndexedArray(
ak.index.Index64([0, 1, 2, 2]),
ak.contents.RecordArray(
[
ak.contents.NumpyArray(
np.array([1, 2, 1], dtype=np.int64)
),
ak.contents.NumpyArray(
np.array([1.1, 2.1, 1.0], dtype=np.float64)
),
],
["pt", "eta"],
),
),
ak.contents.IndexedArray(
ak.index.Index64([0, 0, 1, 2]),
ak.contents.RecordArray(
[
ak.contents.NumpyArray(
np.array([2.1, 2.21, 2.22], dtype=np.float64)
),
ak.contents.NumpyArray(
np.array([2.11, 2.221, 2.222], dtype=np.float64)
),
],
["pt", "eta"],
),
),
],
["foo", "bar"],
),
),
)
)
|
