Skip to content

Commit 1b58ad1

Browse files
authored
Optimize isnull, notnull (#431)
1 parent a7423e6 commit 1b58ad1

File tree

2 files changed

+4
-3
lines changed

2 files changed

+4
-3
lines changed

flox/aggregate_flox.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,8 @@ def ffill(group_idx, array, *, axis, **kwargs):
261261
(group_starts,) = flag.nonzero()
262262

263263
# https://stackoverflow.com/questions/41190852/most-efficient-way-to-forward-fill-nan-values-in-numpy-array
264-
mask = isnull(array)
264+
mask = isnull(array).copy()
265+
# copy needed since we might have a broadcast-trick array
265266
# modified from the SO answer, just reset the index at the start of every group!
266267
mask[..., np.asarray(group_starts)] = False
267268

flox/xrutils.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ def notnull(data):
152152
scalar_type = data.dtype.type
153153
if issubclass(scalar_type, np.bool_ | np.integer | np.character | np.void):
154154
# these types cannot represent missing values
155-
return np.ones_like(data, dtype=bool)
155+
return np.broadcast_to(np.array(True), data.shape)
156156
else:
157157
out = isnull(data)
158158
np.logical_not(out, out=out)
@@ -173,7 +173,7 @@ def isnull(data):
173173
return np.isnan(data)
174174
elif issubclass(scalar_type, np.bool_ | np.integer | np.character | np.void):
175175
# these types cannot represent missing values
176-
return np.zeros_like(data, dtype=bool)
176+
return np.broadcast_to(np.array(False), data.shape)
177177
else:
178178
# at this point, array should have dtype=object
179179
if isinstance(data, (np.ndarray, dask_array_type)): # noqa

0 commit comments

Comments
 (0)