Pickle da.argwhere and da.count_nonzero (#10885)

crusaderky · web-flow · commit c2c1ece0c849 · 2024-02-06T00:23:28.000+01:00
diff --git a/dask/array/routines.py b/dask/array/routines.py
@@ -2056,24 +2056,19 @@ def _isnonzero_vec(v):
 _isnonzero_vec = np.vectorize(_isnonzero_vec, otypes=[bool])
 
 
+def _isnonzero(a):
+    # Output of np.vectorize can't be pickled
+    return _isnonzero_vec(a)
+
+
 def isnonzero(a):
-    if a.dtype.kind in {"U", "S"}:
-        # NumPy treats all-whitespace strings as falsy (like in `np.nonzero`).
-        # but not in `.astype(bool)`. To match the behavior of numpy at least until
-        # 1.19, we use `_isnonzero_vec`. When NumPy changes behavior, we should just
-        # use the try block below.
-        # https://github.com/numpy/numpy/issues/9875
-        return a.map_blocks(_isnonzero_vec, dtype=bool)
+    """Handle special cases where conversion to bool does not work correctly.
+    xref: https://github.com/numpy/numpy/issues/9479
+    """
     try:
-        np.zeros(tuple(), dtype=a.dtype).astype(bool)
+        np.zeros([], dtype=a.dtype).astype(bool)
     except ValueError:
-        ######################################################
-        # Handle special cases where conversion to bool does #
-        # not work correctly.                                #
-        #                                                    #
-        # xref: https://github.com/numpy/numpy/issues/9479   #
-        ######################################################
-        return a.map_blocks(_isnonzero_vec, dtype=bool)
+        return a.map_blocks(_isnonzero, dtype=bool)
     else:
         return a.astype(bool)
 
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
@@ -2,6 +2,7 @@
 
 import contextlib
 import itertools
+import pickle
 import sys
 import warnings
 from numbers import Number
@@ -2708,3 +2709,18 @@ def test_tril_triu_indices(n, k, m, chunks):
         )
     else:
         assert_eq(actual, expected)
+
+
+def test_pickle_vectorized_routines():
+    """Test that graphs that internally use np.vectorize can be pickled"""
+    a = da.from_array(["foo", "bar", ""])
+
+    b = da.count_nonzero(a)
+    assert_eq(b, 2, check_dtype=False)
+    b2 = pickle.loads(pickle.dumps(b))
+    assert_eq(b2, 2, check_dtype=False)
+
+    c = da.argwhere(a)
+    assert_eq(c, [[0], [1]], check_dtype=False)
+    c2 = pickle.loads(pickle.dumps(c))
+    assert_eq(c2, [[0], [1]], check_dtype=False)
