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
|
# 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")
compiler = ROOT.gInterpreter.Declare
# ---------------------------------------------------------------------
def add_column(rdf: ROOT.RDataFrame, arr_val: np.ndarray, name: str):
v_col_org = rdf.GetColumnNames()
l_col_org = [name.c_str() for name in v_col_org]
l_col = []
for col in l_col_org:
l_col.append(col)
data = ak.from_rdataframe(rdf, columns=l_col)
d_data = {col: data[col] for col in l_col}
d_data[name] = arr_val
rdf = ak.to_rdataframe(d_data)
return rdf
# ---------------------------------------------------------------------
def test_add_column():
ak_array_x = ak.Array([1, 2, 3])
ak_array_y = ak.Array([4, 5, 6])
rdf = ak.to_rdataframe({"x": ak_array_x, "y": ak_array_y})
rdf = rdf.Define("z", "ROOT::RVec<int>({1, 2, 3})")
rdf = rdf.Define("w", "true")
arr_val = ak.Array([10, 20, 30])
rdf = add_column(rdf, arr_val, "values")
ROOT.ROOT.DisableImplicitMT()
rdf.Display().Print()
# ---------------------------------------------------------------------
def add_numpy_column(rdf: ROOT.RDataFrame, arr_val: np.ndarray, name: str):
v_col_org = rdf.GetColumnNames()
l_col_org = [name.c_str() for name in v_col_org]
l_col = []
for col in l_col_org:
l_col.append(col)
data = ak.from_rdataframe(rdf, columns=l_col)
d_data = {col: data[col] for col in l_col}
d_data[name] = arr_val
rdf = ak.to_rdataframe(d_data)
return rdf
# ---------------------------------------------------------------------
def test_add_numpy():
d_data = {
"x": np.array([1, 2, 3]),
"y": np.array([4, 5, 6]),
}
rdf = ROOT.RDF.FromNumpy(d_data)
rdf = rdf.Define("z", "ROOT::RVec<int>({1, 2, 3})")
rdf = rdf.Define("w", "true")
arr_val = np.array([10, 20, 30])
rdf = add_numpy_column(rdf, arr_val, "values")
ROOT.ROOT.DisableImplicitMT()
rdf.Display().Print()
|
