nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: dask/dask
Pull request #10318 (2023-05-26 10:29:37)

Nested annotations collide with config 
Fix a bug where, if anything tries reading the dask config from inside two layers of `dask.annotate` context managers, the config may not be there at all!

e.g.
```python
with dask.annotate(optimization=123):
    with dask.annotate(foo=1):
        dask.config.get("optimization.fuse.active")  # TypeError
```


Patch:
diff --git a/dask/base.py b/dask/base.py
index 0743101fc..c9e0491b7 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -153,13 +153,7 @@ def annotate(**annotations):
             % annotations["allow_other_workers"]
         )
 
-    prev_annotations = config.get("annotations", {})
-    new_annotations = {
-        **prev_annotations,
-        **{f"annotations.{k}": v for k, v in annotations.items()},
-    }
-
-    with config.set(new_annotations):
+    with config.set({f"annotations.{k}": v for k, v in annotations.items()}):
         yield
 
 


Test patch:
diff --git a/dask/tests/test_highgraph.py b/dask/tests/test_highgraph.py
index 8ac4ef676..0c69089ef 100644
--- a/dask/tests/test_highgraph.py
+++ b/dask/tests/test_highgraph.py
@@ -182,6 +182,16 @@ def test_multiple_annotations():
     assert clayer.annotations is None
 
 
+def test_annotation_and_config_collision():
+    with dask.config.set({"foo": 1}):
+        with dask.annotate(foo=2):
+            assert dask.config.get("foo") == 1
+            assert dask.config.get("annotations") == {"foo": 2}
+            with dask.annotate(bar=3):
+                assert dask.config.get("foo") == 1
+                assert dask.config.get("annotations") == {"foo": 2, "bar": 3}
+
+
 def test_materializedlayer_cull_preserves_annotations():
     layer = MaterializedLayer(
         {"a": 42, "b": 3.14},
Repository: dask/dask
Pull request #10341 (2023-06-10 00:33:48)

Fix for `total_mem_usage` of `bag.map`
- [x] Closes https://github.com/dask/dask/issues/10338
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`

Relaxing the check to `Iterable` fixes the problem. `itertools.chain` is also an `Iterable`.



Patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 90836281e..3fc6c4264 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -2607,11 +2607,10 @@ def repartition_npartitions(bag, npartitions):
 
 def total_mem_usage(partition):
     from copy import deepcopy
-    from itertools import chain
 
     # if repartition is called multiple times prior to calling compute(), the partitions
-    # will be itertools.chain objects. Copy the object to avoid consuming the iterable.
-    if isinstance(partition, chain):
+    # will be an Iterable. Copy the object to avoid consuming the iterable.
+    if isinstance(partition, Iterable):
         partition = reify(deepcopy(partition))
     return sizeof(partition)
 


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index 1a560f9a1..ab46224b5 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -1681,3 +1681,12 @@ def test_npartitions_saturation(nworkers):
         assert (
             len(db.from_sequence(range(nitems), npartitions=nworkers).dask) >= nworkers
         )
+
+
+def test_map_total_mem_usage():
+    """https://github.com/dask/dask/issues/10338"""
+    b = db.from_sequence(range(1, 100), npartitions=3)
+    total_mem_b = sum(b.map_partitions(total_mem_usage).compute())
+    c = b.map(lambda x: x)
+    total_mem_c = sum(c.map_partitions(total_mem_usage).compute())
+    assert total_mem_b == total_mem_c
Repository: dask/dask
Pull request #8060 (2021-08-18 10:25:19)

Upstream typename instance fix from dask/distributed#5210
In #8019 and dask/distributed#5205 the `typename` util has been consolidated and removed from `distributed` in favour of the `dask` one. Both packages had slightly different implementations of the method.

It looks like `distributed.utils.typename` changed between these two PRs being merged (dask/distributed#5210).

This PR upstreams those changes too and unblocks dask/distributed#5205.

cc @fjetter 


Patch:
diff --git a/dask/utils.py b/dask/utils.py
index 841a5ce6f..7fcb5cb2a 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -826,6 +826,8 @@ def typename(typ, short=False):
     >>> typename(literal, short=True)
     'dask.literal'
     """
+    if not isinstance(typ, type):
+        return typename(type(typ))
     try:
         if not typ.__module__ or typ.__module__ == "builtins":
             return typ.__name__


Test patch:
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
index bd013b63b..5b71881a0 100644
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -724,3 +724,12 @@ def test_noop_context_deprecated():
 def test_typename():
     assert typename(HighLevelGraph) == "dask.highlevelgraph.HighLevelGraph"
     assert typename(HighLevelGraph, short=True) == "dask.HighLevelGraph"
+
+
+class MyType:
+    pass
+
+
+def test_typename_on_instances():
+    instance = MyType()
+    assert typename(instance) == typename(MyType)
Repository: dask/dask
Pull request #10766 (2024-01-10 09:59:49)

Cosmetic: clean up config switches in CI


Patch:
diff --git a/dask/core.py b/dask/core.py
index ef75ca563..e00e4fb8a 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -218,9 +218,26 @@ def iskey(key: object) -> bool:
 
 
 def validate_key(key: object) -> None:
-    """Validate the format of a dask key."""
-    if not iskey(key):
-        raise TypeError(f"Unexpected key type {type(key)} (value: {key!r})")
+    """Validate the format of a dask key.
+
+    See Also
+    --------
+    iskey
+    """
+    if iskey(key):
+        return
+    typ = type(key)
+
+    if typ is tuple:
+        index = None
+        try:
+            for index, part in enumerate(cast(tuple, key)):  # noqa: B007
+                validate_key(part)
+        except TypeError as e:
+            raise TypeError(
+                f"Composite key contains unexpected key type at {index=} ({key=!r})"
+            ) from e
+    raise TypeError(f"Unexpected key type {typ} ({key=!r})")
 
 
 @overload


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index e85973240..69babba06 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -79,17 +79,11 @@ jobs:
 
       - name: Turn on Pyarrow strings config option
         if: ${{ matrix.extra == 'pyarrow' }}
-        run: |
-          export DASK_DATAFRAME__CONVERT_STRING="True"
-          echo "DASK_DATAFRAME__CONVERT_STRING: $DASK_DATAFRAME__CONVERT_STRING"
-          echo "DASK_DATAFRAME__CONVERT_STRING=$DASK_DATAFRAME__CONVERT_STRING" >> $GITHUB_ENV
+        run: echo DASK_DATAFRAME__CONVERT_STRING=True >> $GITHUB_ENV
 
       - name: Turn on dask-expr config option
         if: ${{ matrix.extra == 'dask-expr' }}
-        run: |
-          export DASK_DATAFRAME__QUERY_PLANNING="True"
-          echo "DASK_DATAFRAME__QUERY_PLANNING: DASK_DATAFRAME__QUERY_PLANNING"
-          echo "DASK_DATAFRAME__QUERY_PLANNING=DASK_DATAFRAME__QUERY_PLANNING" >> $GITHUB_ENV
+        run: echo DASK_DATAFRAME__QUERY_PLANNING=True >> $GITHUB_ENV
 
       - name: Reconfigure pytest-timeout
         # No SIGALRM available on Windows
diff --git a/dask/tests/test_core.py b/dask/tests/test_core.py
index a3d72380d..41b42b44d 100644
--- a/dask/tests/test_core.py
+++ b/dask/tests/test_core.py
@@ -73,9 +73,12 @@ def test_iskey_numpy_types():
 def test_validate_key():
     validate_key(1)
     validate_key(("x", 1))
-    with pytest.raises(TypeError):
+    with pytest.raises(TypeError, match="Unexpected key type.*list"):
         validate_key(["x", 1])
 
+    with pytest.raises(TypeError, match="unexpected key type at index=1"):
+        validate_key((2, int))
+
 
 def test_istask():
     assert istask((inc, 1))
Repository: dask/dask
Pull request #10874 (2024-01-29 17:23:15)

Deterministic tokenize for pyarrow datatypes
Fix regression introduced in #10865 (dask-expr specific)


Patch:
diff --git a/dask/base.py b/dask/base.py
index 712e0ba56..1cdbdb28f 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1289,6 +1289,15 @@ def register_pandas():
         return offset.freqstr
 
 
+@normalize_token.register_lazy("pyarrow")
+def register_pyarrow():
+    import pyarrow as pa
+
+    @normalize_token.register(pa.DataType)
+    def normalize_datatype(dt):
+        return pickle.dumps(dt, protocol=4)
+
+
 @normalize_token.register_lazy("numpy")
 def register_numpy():
     import numpy as np


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index d077b8a02..4c26f105c 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -25,6 +25,7 @@ from dask.utils_test import import_or_none
 dd = import_or_none("dask.dataframe")
 np = import_or_none("numpy")
 sp = import_or_none("scipy.sparse")
+pa = import_or_none("pyarrow")
 pd = import_or_none("pandas")
 
 
@@ -781,3 +782,19 @@ def test_tokenize_random_functions_with_state_numpy():
     c = np.random.RandomState(456).random
     assert tokenize(a) == tokenize(b)
     assert tokenize(a) != tokenize(c)
+
+
+@pytest.mark.skipif("not pa")
+def test_tokenize_pyarrow_datatypes_simple():
+    a = pa.int64()
+    b = pa.float64()
+    assert tokenize(a) == tokenize(a)
+    assert tokenize(a) != tokenize(b)
+
+
+@pytest.mark.skipif("not pa")
+def test_tokenize_pyarrow_datatypes_complex():
+    a = pa.struct({"x": pa.int32(), "y": pa.string()})
+    b = pa.struct({"x": pa.float64(), "y": pa.int16()})
+    assert tokenize(a) == tokenize(a)
+    assert tokenize(a) != tokenize(b)
Repository: dask/dask
Pull request #10885 (2024-02-01 19:58:48)

Pickle da.argwhere and da.count_nonzero
Fix serialization of da.argwhere and da.count_nonzero:

```
_pickle.PicklingError: Can't pickle <ufunc '_isnonzero_vec (vectorized)'>: attribute lookup _isnonzero_vec (vectorized) on __main__ failed
```


Patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 8a5c0f922..0398567a0 100644
--- 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)
 


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 5a0e6f285..995f1f7c6 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 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)