Merge branch 'main' into upstream_deps

dask · Feb 7, 2024 · 2d32579 · 2d32579
2 parents ccab90b + 7e04ff7
commit 2d32579
Show file tree

Hide file tree

Showing 19 changed files with 615 additions and 470 deletions.
diff --git a/.github/workflows/update-gpuci.yml b/.github/workflows/update-gpuci.yml
@@ -54,7 +54,7 @@ jobs:
           regex: false
 
       - name: Create Pull Request
-        uses: peter-evans/create-pull-request@v5
+        uses: peter-evans/create-pull-request@v6
         if: ${{ env.UCX_PY_VER != env.NEW_UCX_PY_VER }}  # make sure new ucx-py nightlies are available
         with:
           token: ${{ secrets.GITHUB_TOKEN }}

diff --git a/continuous_integration/environment-3.10.yaml b/continuous_integration/environment-3.10.yaml
@@ -23,7 +23,7 @@ dependencies:
   - pytest-xdist
   - moto<5
   # Optional dependencies
-  - mimesis<13.1.0  # https://github.com/dask/dask/issues/10858
+  - mimesis
   - numpy=1.23
   - pandas=1.5
   - numba<0.59.0  # Removed numba.generated_jit which pandas<2.1 uses

diff --git a/continuous_integration/environment-3.11.yaml b/continuous_integration/environment-3.11.yaml
@@ -23,7 +23,7 @@ dependencies:
   - pytest-xdist
   - moto<5
   # Optional dependencies
-  - mimesis<13.1.0  # https://github.com/dask/dask/issues/10858
+  - mimesis
   - numpy
   - pandas
   - flask

diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
@@ -23,7 +23,7 @@ dependencies:
   - pytest-xdist
   - moto<5
   # Optional dependencies
-  - mimesis<13.1.0  # https://github.com/dask/dask/issues/10858
+  - mimesis
   - numpy
   - pandas
   - flask

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)
diff --git a/dask/array/ufunc.py b/dask/array/ufunc.py
@@ -7,7 +7,7 @@
 
 from dask import core
 from dask.array.core import Array, apply_infer_dtype, asarray, blockwise, elemwise
-from dask.base import is_dask_collection, normalize_function
+from dask.base import is_dask_collection, normalize_token
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils import derived_from, funcname
 
@@ -42,7 +42,7 @@ def __repr__(self):
         return "da.frompyfunc<%s, %d, %d>" % (self._name, self.nin, self.nout)
 
     def __dask_tokenize__(self):
-        return (normalize_function(self._func), self.nin, self.nout)
+        return (normalize_token(self._func), self.nin, self.nout)
 
     def __reduce__(self):
         return (da_frompyfunc, (self._func, self.nin, self.nout))