Skip to content

Commit 65e736e

Browse files
iannaianna
authored
feat: define conversion rules (#59)
* feat: add a method to convert `PyIterable` * feat: add auto conversion rules * fix: update conversion rules * feat: define __init__ function * fix: move convert out of the __init__ function * test: add a test passing a Python array to a Julia function * feat: define conversion rules * fix: remove convertion of an iterable * feat: add rules and register * fix: leave it to a user to invoke pyconvert * test: test pyconvert of an awkward array * tests: add tests * fix: use Vector in convert to avoid problems with BitMasked buffers * feat: convert to an awkward type * fix: add dtypes * fix: more dtypes
1 parent bf7e319 commit 65e736e

File tree

5 files changed

+292
-46
lines changed

5 files changed

+292
-46
lines changed

‎Project.toml

+5-9
Original file line numberDiff line numberDiff line change
@@ -4,23 +4,19 @@ authors = ["Jim Pivarski <[email protected]>", "Jerry Ling <jerry.ling@cern
44
version = "0.1.2"
55

66
[deps]
7+
Conda = "8f4d0f93-b110-5947-807f-2305c1781a2d"
78
JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
8-
Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
9-
10-
[weakdeps]
119
PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
12-
13-
[extensions]
14-
AwkwardPythonCallExt = "PythonCall"
10+
Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
1511

1612
[compat]
1713
JSON = "0.21.4"
18-
julia = "1.9"
1914
Tables = "1.11.1"
15+
julia = "1.9"
16+
PythonCall = "0.9"
2017

2118
[extras]
22-
PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
2319
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
2420

2521
[targets]
26-
test = ["Test", "PythonCall"]
22+
test = ["Test"]

‎ext/AwkwardPythonCallExt/AwkwardPythonCallExt.jl

-35
This file was deleted.

‎src/AwkwardArray.jl

+3-2
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@ import Tables
66
include("./all_implementations.jl")
77
include("./tables.jl")
88

9-
# stub for PythonCall Extention
10-
function convert end
9+
include("./AwkwardPythonCallExt.jl")
10+
using .AwkwardPythonCallExt: convert
11+
1112

1213
end # module AwkwardArray

‎src/AwkwardPythonCallExt.jl

+116
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
module AwkwardPythonCallExt
2+
using PythonCall
3+
using JSON
4+
import AwkwardArray
5+
6+
function AwkwardArray.convert(layout::AwkwardArray.Content)::Py
7+
form, len, containers = AwkwardArray.to_buffers(layout)
8+
9+
py_buffers = Dict{String,Any}()
10+
11+
for (key, buffer) in containers
12+
py_buffers[key] = pyimport("numpy").asarray(buffer, dtype = pyimport("numpy").uint8)
13+
end
14+
15+
pyimport("awkward").from_buffers(form, len, py_buffers)
16+
end
17+
18+
function AwkwardArray.convert(array::Py)::AwkwardArray.Content
19+
form, len, _containers = pyimport("awkward").to_buffers(array)
20+
containers = pyconvert(Dict, _containers)
21+
22+
julia_buffers = Dict{String,Vector{UInt8}}()
23+
24+
for (key, buffer) in containers
25+
julia_buffers[key] = reinterpret(UInt8, buffer)
26+
end
27+
28+
AwkwardArray.from_buffers(
29+
pyconvert(String, form.to_json()),
30+
pyconvert(Int, len),
31+
julia_buffers,
32+
)
33+
end
34+
35+
# rule functions
36+
function pyconvert_rule_awkward_array_primitive(::Type{AwkwardArray.PrimitiveArray}, x::Py)
37+
array = AwkwardArray.convert(x)
38+
return PythonCall.pyconvert_return(array)
39+
end
40+
41+
function pyconvert_rule_awkward_array_empty(::Type{AwkwardArray.EmptyArray}, x::Py)
42+
array = AwkwardArray.convert(x)
43+
return PythonCall.pyconvert_return(array)
44+
end
45+
46+
function pyconvert_rule_awkward_array_listoffset(::Type{AwkwardArray.ListOffsetArray}, x::Py)
47+
array = AwkwardArray.convert(x)
48+
return PythonCall.pyconvert_return(array)
49+
end
50+
51+
function pyconvert_rule_awkward_array_list(::Type{AwkwardArray.ListArray}, x::Py)
52+
array = AwkwardArray.convert(x)
53+
return PythonCall.pyconvert_return(array)
54+
end
55+
function pyconvert_rule_awkward_array_regular(::Type{AwkwardArray.RegularArray}, x::Py)
56+
array = AwkwardArray.convert(x)
57+
return PythonCall.pyconvert_return(array)
58+
end
59+
60+
function pyconvert_rule_awkward_array_record(::Type{AwkwardArray.RecordArray}, x::Py)
61+
array = AwkwardArray.convert(x)
62+
return PythonCall.pyconvert_return(array)
63+
end
64+
65+
function pyconvert_rule_awkward_array_tuple(::Type{AwkwardArray.TupleArray}, x::Py)
66+
array = AwkwardArray.convert(x)
67+
return PythonCall.pyconvert_return(array)
68+
end
69+
70+
function pyconvert_rule_awkward_array_indexed(::Type{AwkwardArray.IndexedArray}, x::Py)
71+
array = AwkwardArray.convert(x)
72+
return PythonCall.pyconvert_return(array)
73+
end
74+
75+
function pyconvert_rule_awkward_array_indexedoption(::Type{AwkwardArray.IndexedOptionArray}, x::Py)
76+
array = AwkwardArray.convert(x)
77+
return PythonCall.pyconvert_return(array)
78+
end
79+
80+
function pyconvert_rule_awkward_array_bytemasked(::Type{AwkwardArray.ByteMaskedArray}, x::Py)
81+
array = AwkwardArray.convert(x)
82+
return PythonCall.pyconvert_return(array)
83+
end
84+
85+
function pyconvert_rule_awkward_array_bitmasked(::Type{AwkwardArray.BitMaskedArray}, x::Py)
86+
array = AwkwardArray.convert(x)
87+
return PythonCall.pyconvert_return(array)
88+
end
89+
90+
function pyconvert_rule_awkward_array_unmasked(::Type{AwkwardArray.UnmaskedArray}, x::Py)
91+
array = AwkwardArray.convert(x)
92+
return PythonCall.pyconvert_return(array)
93+
end
94+
95+
function pyconvert_rule_awkward_array_union(::Type{AwkwardArray.UnionArray}, x::Py)
96+
array = AwkwardArray.convert(x)
97+
return PythonCall.pyconvert_return(array)
98+
end
99+
100+
function __init__()
101+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.PrimitiveArray, pyconvert_rule_awkward_array_primitive, PythonCall.PYCONVERT_PRIORITY_ARRAY)
102+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.EmptyArray, pyconvert_rule_awkward_array_empty, PythonCall.PYCONVERT_PRIORITY_ARRAY)
103+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.ListOffsetArray, pyconvert_rule_awkward_array_listoffset, PythonCall.PYCONVERT_PRIORITY_ARRAY)
104+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.ListArray, pyconvert_rule_awkward_array_list, PythonCall.PYCONVERT_PRIORITY_ARRAY)
105+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.RegularArray, pyconvert_rule_awkward_array_regular, PythonCall.PYCONVERT_PRIORITY_ARRAY)
106+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.RecordArray, pyconvert_rule_awkward_array_record, PythonCall.PYCONVERT_PRIORITY_ARRAY)
107+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.TupleArray, pyconvert_rule_awkward_array_tuple, PythonCall.PYCONVERT_PRIORITY_ARRAY)
108+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.IndexedArray, pyconvert_rule_awkward_array_indexed, PythonCall.PYCONVERT_PRIORITY_ARRAY)
109+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.IndexedOptionArray, pyconvert_rule_awkward_array_indexedoption, PythonCall.PYCONVERT_PRIORITY_ARRAY)
110+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.ByteMaskedArray, pyconvert_rule_awkward_array_bytemasked, PythonCall.PYCONVERT_PRIORITY_ARRAY)
111+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.BitMaskedArray, pyconvert_rule_awkward_array_bitmasked, PythonCall.PYCONVERT_PRIORITY_ARRAY)
112+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.UnmaskedArray, pyconvert_rule_awkward_array_unmasked, PythonCall.PYCONVERT_PRIORITY_ARRAY)
113+
PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.UnionArray, pyconvert_rule_awkward_array_union, PythonCall.PYCONVERT_PRIORITY_ARRAY)
114+
end
115+
116+
end # module

‎test/runpytests.jl

+168
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,171 @@ end
3030

3131
@test array == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
3232
end
33+
34+
# Test pyconvert Python Awkwar Array to Julia Awkward Array
35+
@testset "convert # PrimitiveArray" begin
36+
layout = pyimport("awkward").contents.NumpyArray(
37+
pyimport("numpy").array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], dtype=pyimport("numpy").float64)
38+
)
39+
py_array = pyimport("awkward").Array(layout)
40+
41+
array = pyconvert(AwkwardArray.PrimitiveArray, py_array)
42+
@test array isa AwkwardArray.PrimitiveArray
43+
end
44+
45+
@testset "convert # EmptyArray" begin
46+
layout = pyimport("awkward").contents.EmptyArray()
47+
py_array = pyimport("awkward").Array(layout)
48+
49+
array = pyconvert(AwkwardArray.EmptyArray, py_array)
50+
@test array isa AwkwardArray.EmptyArray
51+
end
52+
53+
@testset "convert # ListOffsetArray" begin
54+
py_array = pyimport("awkward").Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
55+
56+
array = pyconvert(AwkwardArray.ListOffsetArray, py_array)
57+
@test array isa AwkwardArray.ListOffsetArray
58+
end
59+
60+
@testset "convert # ListArray" begin
61+
content = pyimport("awkward").contents.NumpyArray(
62+
pyimport("numpy").array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], dtype=pyimport("numpy").float64)
63+
)
64+
starts = pyimport("awkward").index.Index64(pyimport("numpy").array([0, 3, 3, 5, 6], dtype=pyimport("numpy").int64))
65+
stops = pyimport("awkward").index.Index64(pyimport("numpy").array([3, 3, 5, 6, 9], dtype=pyimport("numpy").int64))
66+
offsets = pyimport("awkward").index.Index64(pyimport("numpy").array([0, 3, 3, 5, 6, 9], dtype=pyimport("numpy").int64))
67+
layout = pyimport("awkward").contents.ListArray(starts, stops, content)
68+
69+
py_array = pyimport("awkward").Array(layout)
70+
71+
array = pyconvert(AwkwardArray.ListArray, py_array)
72+
@test array isa AwkwardArray.ListArray
73+
end
74+
75+
@testset "convert # RegularArray" begin
76+
content = pyimport("awkward").contents.NumpyArray(
77+
pyimport("numpy").array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], dtype=pyimport("numpy").float64)
78+
)
79+
offsets = pyimport("awkward").index.Index64(pyimport("numpy").array([0, 3, 3, 5, 6, 10, 10], dtype=pyimport("numpy").int64))
80+
listoffsetarray = pyimport("awkward").contents.ListOffsetArray(offsets, content)
81+
regulararray = pyimport("awkward").contents.RegularArray(listoffsetarray, 2, zeros_length=0)
82+
83+
py_array = pyimport("awkward").Array(regulararray)
84+
85+
array = pyconvert(AwkwardArray.RegularArray, py_array)
86+
@test array isa AwkwardArray.RegularArray
87+
end
88+
89+
@testset "convert # RecordArray" begin
90+
content1 = pyimport("awkward").contents.NumpyArray(pyimport("numpy").array([1, 2, 3, 4, 5], dtype=pyimport("numpy").int64))
91+
content2 = pyimport("awkward").contents.NumpyArray(
92+
pyimport("numpy").array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], dtype=pyimport("numpy").float64)
93+
)
94+
offsets = pyimport("awkward").index.Index64(pyimport("numpy").array([0, 3, 3, 5, 6, 9], dtype=pyimport("numpy").int64))
95+
listoffsetarray = pyimport("awkward").contents.ListOffsetArray(offsets, content2)
96+
recordarray = pyimport("awkward").contents.RecordArray(
97+
[content1, listoffsetarray, content2, content1],
98+
fields=["one", "two", "2", "wonky"],
99+
)
100+
101+
py_array = pyimport("awkward").Array(recordarray)
102+
103+
array = pyconvert(AwkwardArray.RecordArray, py_array)
104+
@test array isa AwkwardArray.RecordArray
105+
end
106+
107+
@testset "convert # TupleArray" begin
108+
tuplearray = pyimport("awkward").contents.RecordArray([pyimport("awkward").contents.NumpyArray(pyimport("numpy").arange(10, dtype=pyimport("numpy").int64))], pybuiltins.None)
109+
110+
py_array = pyimport("awkward").Array(tuplearray)
111+
112+
array = pyconvert(AwkwardArray.TupleArray, py_array)
113+
@test array isa AwkwardArray.TupleArray
114+
end
115+
116+
@testset "convert # IndexedArray" begin
117+
content = pyimport("awkward").contents.NumpyArray(pyimport("numpy").array([0.0, 1.1, 2.2, 3.3, 4.4], dtype=pyimport("numpy").float64))
118+
119+
ind = pyimport("numpy").array([2, 2, 0, 3, 4], dtype=pyimport("numpy").int32)
120+
index = pyimport("awkward").index.Index32(ind)
121+
indexedarray = pyimport("awkward").contents.IndexedArray(index, content)
122+
123+
py_array = pyimport("awkward").Array(indexedarray)
124+
125+
array = pyconvert(AwkwardArray.IndexedArray, py_array)
126+
@test array isa AwkwardArray.IndexedArray
127+
end
128+
129+
@testset "convert # IndexedOptionArray" begin
130+
content = pyimport("awkward").contents.NumpyArray(pyimport("numpy").array([0.0, 1.1, 2.2, 3.3, 4.4], dtype=pyimport("numpy").float64))
131+
index = pyimport("awkward").index.Index64(pyimport("numpy").array([2, 2, 0, -1, 4], dtype=pyimport("numpy").int64))
132+
indexedoptionarray = pyimport("awkward").contents.IndexedOptionArray(index, content)
133+
134+
py_array = pyimport("awkward").Array(indexedoptionarray)
135+
136+
array = pyconvert(AwkwardArray.IndexedOptionArray, py_array)
137+
@test array isa AwkwardArray.IndexedOptionArray
138+
end
139+
140+
@testset "convert # ByteMaskedArray" begin
141+
layout = pyimport("awkward").contents.ByteMaskedArray(
142+
pyimport("awkward").index.Index8(pyimport("numpy").array([0, 1, 0, 1, 0], dtype=pyimport("numpy").int8)),
143+
pyimport("awkward").contents.NumpyArray(pyimport("numpy").arange(5, dtype=pyimport("numpy").int64)),
144+
valid_when=pybuiltins.True,
145+
)
146+
py_array = pyimport("awkward").Array(layout)
147+
148+
array = pyconvert(AwkwardArray.ByteMaskedArray, py_array)
149+
@test array isa AwkwardArray.ByteMaskedArray
150+
end
151+
152+
@testset "convert # BitMaskedArray" begin
153+
content = pyimport("awkward").operations.from_iter(
154+
[[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=pybuiltins.False
155+
)
156+
mask = pyimport("awkward").index.IndexU8(pyimport("numpy").array([66], dtype=pyimport("numpy").uint8))
157+
maskedarray = pyimport("awkward").contents.BitMaskedArray(
158+
mask, content, valid_when=pybuiltins.False, length=4, lsb_order=pybuiltins.True
159+
)
160+
py_array = pyimport("awkward").Array(maskedarray)
161+
162+
array = pyconvert(AwkwardArray.BitMaskedArray, py_array)
163+
@test array isa AwkwardArray.BitMaskedArray
164+
end
165+
166+
@testset "convert # UnmaskedArray" begin
167+
unmaskedarray = pyimport("awkward").contents.UnmaskedArray(
168+
pyimport("awkward").contents.NumpyArray(
169+
pyimport("numpy").array([0.0, 1.1, 2.2, 3.3], dtype=pyimport("numpy").float64)
170+
)
171+
)
172+
py_array = pyimport("awkward").Array(unmaskedarray)
173+
174+
array = pyconvert(AwkwardArray.UnmaskedArray, py_array)
175+
@test array isa AwkwardArray.UnmaskedArray
176+
end
177+
178+
@testset "convert # UnionArray" begin
179+
layout = pyimport("awkward").contents.unionarray.UnionArray(
180+
pyimport("awkward").index.Index(pyimport("numpy").array([1, 1, 0, 0, 1, 0, 1], dtype=pyimport("numpy").int8)),
181+
pyimport("awkward").index.Index(pyimport("numpy").array([4, 3, 0, 1, 2, 2, 4, 100], dtype=pyimport("numpy").int64)),
182+
[
183+
pyimport("awkward").contents.recordarray.RecordArray(
184+
[pyimport("awkward").from_iter(["1", "2", "3"], highlevel=pybuiltins.False)], ["nest"]
185+
),
186+
pyimport("awkward").contents.recordarray.RecordArray(
187+
[
188+
pyimport("awkward").contents.numpyarray.NumpyArray(
189+
pyimport("numpy").array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=pyimport("numpy").float64)
190+
)
191+
],
192+
["nest"],
193+
),
194+
],
195+
)
196+
py_array = pyimport("awkward").Array(layout)
197+
198+
array = pyconvert(AwkwardArray.UnionArray, py_array)
199+
@test array isa AwkwardArray.UnionArray
200+
end

0 commit comments

Comments
 (0)