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
95
96
97
98
99
100
101
102
103
104
105
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import os
import numpy as np # noqa: F401
import pytest
import awkward as ak
ROOT = pytest.importorskip("ROOT")
ROOT.ROOT.EnableImplicitMT(1)
compiler = ROOT.gInterpreter.Declare
def test_data_frame_integers(tmp_path):
filename = os.path.join(tmp_path, "test-integers.root")
ak_array_x = ak.Array([1, 2, 3, 4, 5])
ak_array_y = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5])
data_frame = ak.to_rdataframe({"x": ak_array_x, "y": ak_array_y})
assert data_frame.GetColumnType("x") == "int64_t"
assert data_frame.GetColumnType("y") == "double"
ak_array_out = ak.from_rdataframe(
data_frame,
columns=("x", "y"),
)
assert ak_array_x.to_list() == ak_array_out["x"].to_list()
assert ak_array_y.to_list() == ak_array_out["y"].to_list()
data_frame.Snapshot("Test", filename, ("x", "y"))
def test_data_frame_vec_of_vec_of_real(tmp_path):
import warnings
filename = os.path.join(tmp_path, "test-listarray.root")
ak_array_in = ak.Array([[[1.1], [2.2]], [[3.3], [4.4, 5.5]]])
data_frame = ak.to_rdataframe({"x": ak_array_in})
assert data_frame.GetColumnType("x").startswith("awkward::ListArray_")
ak_array_out = ak.from_rdataframe(
data_frame,
columns=("x",),
)
assert ak_array_in.to_list() == ak_array_out["x"].to_list()
# With `ROOT.ROOT.EnableImplicitMT(1)` a SystemError becomes a Warning:
with warnings.catch_warnings(record=True):
# Warning in <TStreamerInfo::Build>: awkward::ListArray_jEomw7jWD1w:
# base class awkward::ArrayView has no streamer or dictionary
# it will not be saved
data_frame.Snapshot("ListArray", filename, ("x",))
def test_data_frame_rvec_filter(tmp_path):
filename = os.path.join(tmp_path, "test-listarray2.root")
ak_array_x = ak.Array([[1, 2], [3], [4, 5]])
ak_array_y = ak.Array([[1.0, 1.1], [2.2, 3.3, 4.4], [5.5]])
data_frame = ak.to_rdataframe({"x": ak_array_x, "y": ak_array_y})
rdf3 = data_frame.Filter("x.size() >= 2")
assert data_frame.GetColumnType("x") == "ROOT::VecOps::RVec<int64_t>"
assert data_frame.GetColumnType("y") == "ROOT::VecOps::RVec<double>"
ak_array_out = ak.from_rdataframe(
rdf3,
columns=(
"x",
"y",
),
)
assert ak_array_out["x"].to_list() == [[1, 2], [4, 5]]
assert ak_array_out["y"].to_list() == [[1.0, 1.1], [5.5]]
rdf4 = data_frame.Filter("y.size() == 2")
ak_array_out = ak.from_rdataframe(
rdf4,
columns=(
"x",
"y",
),
)
assert ak_array_out["x"].to_list() == [[1, 2]]
assert ak_array_out["y"].to_list() == [[1.0, 1.1]]
data_frame.Snapshot(
"ListArray",
filename,
(
"x",
"y",
),
)
|
