Optimize isnull, notnull (#431)

Author: dcherian

flox/aggregate_flox.py

@@ -261,7 +261,8 @@ def ffill(group_idx, array, *, axis, **kwargs):
     (group_starts,) = flag.nonzero()
 
     # https://stackoverflow.com/questions/41190852/most-efficient-way-to-forward-fill-nan-values-in-numpy-array
-    mask = isnull(array)
+    mask = isnull(array).copy()
+    # copy needed since we might have a broadcast-trick array
     # modified from the SO answer, just reset the index at the start of every group!
     mask[..., np.asarray(group_starts)] = False
 

flox/xrutils.py

@@ -152,7 +152,7 @@ def notnull(data):
     scalar_type = data.dtype.type
     if issubclass(scalar_type, np.bool_ | np.integer | np.character | np.void):
         # these types cannot represent missing values
-        return np.ones_like(data, dtype=bool)
+        return np.broadcast_to(np.array(True), data.shape)
     else:
         out = isnull(data)
         np.logical_not(out, out=out)
@@ -173,7 +173,7 @@ def isnull(data):
         return np.isnan(data)
     elif issubclass(scalar_type, np.bool_ | np.integer | np.character | np.void):
         # these types cannot represent missing values
-        return np.zeros_like(data, dtype=bool)
+        return np.broadcast_to(np.array(False), data.shape)
     else:
         # at this point, array should have dtype=object
         if isinstance(data, (np.ndarray, dask_array_type)):  # noqa