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
|
# 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
numba = pytest.importorskip("numba")
ak.numba.register_and_check()
def test_unmasked():
@numba.njit
def find_it(array):
for item in array:
if item is None:
pass
elif item.x == 3:
return item
return None
content = ak.highlevel.Array([{"x": 1}, {"x": 2}, {"x": 3}]).layout
unmasked = ak.contents.UnmaskedArray(content)
array = ak.highlevel.Array(unmasked)
assert ak.operations.to_list(find_it(array)) == {"x": 3}
def test_indexedoption():
@numba.njit
def find_it(array):
for item in array:
if item is None:
pass
elif item.x == 3:
return item
return None
array = ak.highlevel.Array([{"x": 1}, {"x": 2}, None, {"x": 3}])
assert ak.operations.to_list(find_it(array)) == {"x": 3}
def test_indexed_1():
@numba.njit
def f1(array, check):
for i in range(len(array)):
item = array[i]
if item.x == check:
return i
return 999
content = ak.highlevel.Array([{"x": 100}, {"x": 101}, {"x": 102}]).layout
index = ak.index.Index64(np.array([2, 0, 1], dtype=np.int64))
indexedarray = ak.contents.IndexedArray(index, content)
array = ak.highlevel.Array(indexedarray)
assert f1(array, 100) == 1
assert f1(array, 101) == 2
assert f1(array, 102) == 0
assert f1(array, 12345) == 999
def test_indexed_2():
@numba.njit
def f1(array, check):
for item in array:
if item.x == check:
return item
return None
content = ak.highlevel.Array([{"x": 100}, {"x": 101}, {"x": 102}]).layout
index = ak.index.Index64(np.array([2, 0, 1], dtype=np.int64))
indexedarray = ak.contents.IndexedArray(index, content)
array = ak.highlevel.Array(indexedarray)
assert f1(array, 100).to_list() == {"x": 100}
assert f1(array, 101).to_list() == {"x": 101}
assert f1(array, 102).to_list() == {"x": 102}
assert f1(array, 12345) is None
|
