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
|
# 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
import awkward._connect.cling
import awkward._lookup
ROOT = pytest.importorskip("ROOT")
ROOT.ROOT.EnableImplicitMT(1)
compiler = ROOT.gInterpreter.Declare
def test_to_from_data_frame_large():
n = 6
assert 2 * (n // 2) == n
rows = 3 ** (n // 2)
cols = n
arr = np.zeros((rows, cols), dtype=np.int64)
shape = (rows,)
source = np.array([-1, 0, 1], dtype=np.int64)[:, None]
for col in range(n // 2):
shape = (
-1,
3,
shape[-1] // 3,
)
col_view = arr[:, col]
col_view.shape = shape
col_view[:] = source
ak_array_in = ak.from_numpy(arr, regulararray=True)
data_frame = ak.to_rdataframe({"x": ak_array_in})
ak_array_out = ak.from_rdataframe(
data_frame,
columns=("x",),
highlevel=False,
)
assert isinstance(ak_array_out, ak.contents.RecordArray)
assert len(ak_array_in) == len(ak_array_out)
def test_data_frame_integers():
ak_array_in = ak.Array([1, 2, 3, 4, 5])
data_frame = ak.to_rdataframe({"x": ak_array_in})
assert data_frame.GetColumnType("x") == "int64_t"
def overload_add(left, right):
return ak.Array({"x": left.x + right.x})
behavior = {}
behavior[np.add, "Overload", "Overload"] = overload_add
ak_array_out = ak.from_rdataframe(
data_frame,
columns=("x",),
keep_order=True,
behavior=behavior,
with_name="Overload",
)
assert ak_array_in.to_list() == ak_array_out["x"].to_list()
assert (ak_array_out + ak_array_out).to_list() == [
{"x": 2},
{"x": 4},
{"x": 6},
{"x": 8},
{"x": 10},
]
|
