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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import numpy as np
import awkward as ak
from awkward.operations import to_list
def test_ak_where_with_optional_unknowns():
"""
This is the example from the Issue.
In the two cases we fail, the not-selected value has type ?unknown, value None of course.
Names here are changed a little from github issue 3098.
"""
opt_true_cond = ak.Array([[True], [None]])[0] # <Array [True] type='1 * ?bool'>
true_cond = ak.Array([True]) # <Array [True] type='1 * bool'>
none_alternative = ak.Array([None]) # <Array [None] type='1 * ?unknown'>
zero_alternative = ak.Array([0]) # <Array [0] type='1 * int64'>
opt_zero_alternative = ak.Array([[0], [None]])[0] # <Array [0] type='1 * ?int64'>
assert ak.where(opt_true_cond, 1, none_alternative).to_list() == [1]
# Above fails at time of writing. We get [None].
assert ak.where(opt_true_cond, 1, zero_alternative).to_list() == [1]
assert ak.where(opt_true_cond, 1, opt_zero_alternative).to_list() == [1]
# These assertions pass. Note that true_cond is type bool, not ?bool.
assert ak.where(true_cond, 1, none_alternative).to_list() == [1]
assert ak.where(true_cond, 1, zero_alternative).to_list() == [1]
assert ak.where(true_cond, 1, opt_zero_alternative).to_list() == [1]
# Like the first three assertions, The first one here fails.
# This demonstrates that the problem is symmetric w/rt X and Y arrays.
assert ak.where(~opt_true_cond, none_alternative, 1).to_list() == [1]
# Above fails at time of writing. We're getting [None] again.
assert ak.where(~opt_true_cond, zero_alternative, 1).to_list() == [1]
assert ak.where(~opt_true_cond, opt_zero_alternative, 1).to_list() == [1]
def test_ak_where_with_optionals():
"""
It turns out that we don't need to use ?unknown arrays to trigger this issue.
We only need a None (masked value) in an element that is selected against.
At time-of-writing (ATOW), ak.where() produces None values when:
1. The conditional array values have OptionType (at least ?bool or ?int64), *AND*
2. The value array element *NOT* selected has Option type and holds a None value.
In this case regardless of the type or value of the array element that *IS* selected,
the result for that element will, incorrectly, be None.
"""
# This passes. Note that a condition of None creates a None in the result.
assert to_list(
ak.where(
ak.Array([True, False, None]), ak.Array([1, 2, 3]), ak.Array([4, 5, 6])
)
) == [1, 5, None]
# This also passes. (The presence of None at the end forces option types to be used.)
assert to_list(
ak.where(
ak.Array([True, False, None]),
ak.Array([1, 2, None]),
ak.Array([4, 5, None]),
)
) == [1, 5, None]
# This fails (ATOW). The presence of None forces option types to be used.
assert to_list(
ak.where(
ak.Array([True, False, None]),
ak.Array([1, 2, None]),
ak.Array([None, 5, None]),
)
) == [1, 5, None] # ATOW we get [None, 5, None]
# Fails ATOW. Same as above but with a None in the X argument.
assert to_list(
ak.where(
ak.Array([True, False, None]),
ak.Array([1, None, None]),
ak.Array([4, 5, None]),
)
) == [1, 5, None] # ATOW we get [1, None, None]
# Fails ATOW. Same as above the Y argument is not even an optional type (but X still is).
assert to_list(
ak.where(
ak.Array([True, False, None]),
ak.Array([1, None, None]),
ak.Array([4, 5, 6]),
)
) == [1, 5, None] # ATOW we get [1, None, None]
def test_ak_where_with_optionals_multidim():
# This needs to continue to work:
assert to_list(
ak.where(
ak.Array([True, False]),
ak.Array([[1, 2], [3, 4]]),
ak.Array([[10, 11], [12, 13]]),
)
) == [[1, 2], [12, 13]]
# Option types only in X, not condition
assert to_list(
ak.where(
ak.Array([[True, True], [False, False]]),
ak.Array([[1, 2], None]),
ak.Array([[10, 11], [12, 13]]),
)
) == [[1, 2], [12, 13]]
# Option types in condition and X, only one level of depth
assert to_list(
ak.where(
ak.Array([[True, True], [False, False], [True, False], None]),
ak.Array([[1, 2], None, None, [7, 8]]),
ak.Array([[11, 12], [13, 14], [15, 16], [17, 18]]),
)
) == [[1, 2], [13, 14], [None, 16], [None, None]]
# Note: [[1, 2], [13, 14], [None, 16], None] might seem more natural,
# but broadcasting expands these arrays out.
assert to_list(
ak.where(
ak.Array([[True, False], [True, None]]),
ak.Array([1, 2]),
ak.Array([None, 12]),
)
) == [[1, None], [2, None]]
def test_ak_where_more_option_types():
assert to_list(
ak.where(
ak.Array([False, True, None]),
ak.Array(["this", None, "that"]),
ak.Array(["foo", "bar", "baz"]),
)
) == ["foo", None, None]
bitmasked5 = ak.contents.BitMaskedArray(
mask=ak.index.Index(
np.array(
[
0b10100,
],
dtype=np.uint8,
)
),
content=ak.contents.NumpyArray(np.arange(5)),
valid_when=False,
length=5,
lsb_order=True,
parameters={"_my_param": "boysenberry"},
) # [0, 1, None, 3, None]
unmasked5 = ak.contents.UnmaskedArray(
ak.contents.NumpyArray(np.arange(10, 15))
) # [10, 11, 12, 13, 14]
union5 = ak.Array([True, None, "two", 3, 4.4])
mixed_result = ak.where(
ak.Array([True, None, True, False, True]), bitmasked5, unmasked5
)
assert to_list(mixed_result) == [0, None, None, 13, None]
assert (
mixed_result.layout.parameters.get("_my_param") is None
) # Params not preserved.
assert to_list(
ak.where(ak.Array([True, True, True, False, None]), union5, unmasked5)
) == [True, None, "two", 13, None]
|
