Repository: dask/dask
Pull request #10701 (2023-12-13 17:54:13)

Relax redundant-key check in ``_check_dsk``
Possible solution for the `test_setitem_hardmask` component of https://github.com/dask/dask/issues/10672

It seems that we are knowingly adding redundant keys for array setitem operations (https://github.com/dask/dask/blob/main/dask/array/slicing.py#L2064).  It seems like we are only catching this in Windows, because other systems seem to tokenize distinct `np.ma.masked` objects to different values. I added a new test (`test_setitem_slice_twice`) that also fails on Linux without the proposed change.

**Question**: Should we relax the redundant-key check in `_check_dsk` (as this PR currently does). Or, should we avoid doing things like [this](https://github.com/dask/dask/blob/main/dask/array/slicing.py#L2064)?

cc @charlesbluca 



Patch:
diff --git a/dask/array/utils.py b/dask/array/utils.py
index fd7a0e2a9..39da95be7 100644
--- a/dask/array/utils.py
+++ b/dask/array/utils.py
@@ -12,6 +12,7 @@ from tlz import concat, frequencies
 
 from dask.array.core import Array
 from dask.array.numpy_compat import AxisError
+from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils import has_keyword, is_arraylike, is_cupy_type, typename
 
@@ -210,7 +211,15 @@ def _check_dsk(dsk):
     assert all(isinstance(k, (tuple, str)) for k in dsk.layers)
     freqs = frequencies(concat(dsk.layers.values()))
     non_one = {k: v for k, v in freqs.items() if v != 1}
-    assert not non_one, non_one
+    key_collisions = set()
+    # Allow redundant keys if the values are equivalent
+    for k in non_one.keys():
+        for layer in dsk.layers.values():
+            try:
+                key_collisions.add(tokenize(layer[k]))
+            except KeyError:
+                pass
+    assert len(key_collisions) < 2, non_one
 
 
 def assert_eq_shape(a, b, check_ndim=True, check_nan=True):
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 6c6115bc5..7cddc42d9 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1714,9 +1714,23 @@ Dask Name: {name}, {layers}"""
         pandas.DataFrame.memory_usage
         """
         if isinstance(divisions, int):
+            warnings.warn(
+                "divisions is an integer and will be inferred as npartitions instead. "
+                "This automatic inference is deprecated and will change in the future. "
+                f"Please set npartitions={divisions} instead.",
+                FutureWarning,
+                stacklevel=2,
+            )
             npartitions = divisions
             divisions = None
         if isinstance(divisions, str):
+            warnings.warn(
+                "divisions is a string and will be inferred as partition_size instead. "
+                "This automatic inference is deprecated and will change in the future. "
+                f"Please set partition_size={divisions} instead.",
+                FutureWarning,
+                stacklevel=2,
+            )
             partition_size = divisions
             divisions = None
         if (
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 8500691ac..a0781200c 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -60,6 +60,63 @@ recent changes made to ``dask.order``.
 
 See :pr-distributed:`8415`, :pr-distributed:`8416`, and :pr-distributed:`8414` from `Hendrik Makait`_ for details.
 
+
+Removed disabling pickle option
+"""""""""""""""""""""""""""""""
+
+The ``distributed.scheduler.pickle`` configuration option is no longer supported.
+As of the 2023.4.0 release, ``pickle`` is used to transmit task graphs, so can no
+longer be disabled. We now raise an informative error when ``distributed.scheduler.pickle``
+is set to ``False``.
+
+See :pr-distributed:`8401` from `Florian Jetter`_ for details.
+
+
+.. dropdown:: Additional changes
+
+  - Add changelog entry for recent P2P merge fixes (:pr:`10712`) `Hendrik Makait`_
+  - Update DataFrame page (:pr:`10710`) `Matthew Rocklin`_
+  - Add changelog entry for ``dask-expr`` switch (:pr:`10704`) `Patrick Hoefler`_
+  - Improve changelog entry for ``PipInstall`` changes (:pr:`10711`) `Hendrik Makait`_
+  - Remove PR labeler (:pr:`10709`) `James Bourbeau`_
+  - Add ``.__wrapped__`` to ``Delayed`` object (:pr:`10695`) `Andrew S. Rosen`_
+  - Bump ``actions/labeler`` from 4.3.0 to 5.0.0 (:pr:`10689`)
+  - Bump ``actions/stale`` from 8 to 9 (:pr:`10690`)
+  - [Dask.order] Remove non-runnable leaf nodes from ordering (:pr:`10697`) `Florian Jetter`_
+  - Update installation docs (:pr:`10699`) `Matthew Rocklin`_
+  - Fix software environment link in docs (:pr:`10700`) `James Bourbeau`_
+  - Avoid converting non-strings to arrow strings for read_parquet (:pr:`10692`) `Patrick Hoefler`_
+  - Bump ``xarray-contrib/issue-from-pytest-log`` from 1.2.7 to 1.2.8 (:pr:`10687`)
+  - Fix ``tokenize`` for ``pd.DateOffset`` (:pr:`10664`) `jochenott`_
+  - Bugfix for writing empty array to zarr (:pr:`10506`) `Ben`_
+  - Docs update, fixup styling, mention free (:pr:`10679`) `Matthew Rocklin`_
+  - Update deployment docs (:pr:`10680`) `Matthew Rocklin`_
+  - Dask.order rewrite using a critical path approach (:pr:`10660`) `Florian Jetter`_
+  - Avoid substituting keys that occur multiple times (:pr:`10646`) `Florian Jetter`_
+  - Add missing image to docs (:pr:`10694`) `Matthew Rocklin`_
+  - Bump ``actions/setup-python`` from 4 to 5 (:pr:`10688`)
+  - Update landing page (:pr:`10674`) `Matthew Rocklin`_
+  - Make meta check simpler in dispatch (:pr:`10638`) `Patrick Hoefler`_
+  - Pin PR Labeler (:pr:`10675`) `Matthew Rocklin`_
+  - Reorganize docs index a bit (:pr:`10669`) `Matthew Rocklin`_
+  - Bump ``actions/setup-java`` from 3 to 4 (:pr:`10667`)
+  - Bump ``conda-incubator/setup-miniconda`` from 2.2.0 to 3.0.1 (:pr:`10668`)
+  - Bump ``xarray-contrib/issue-from-pytest-log`` from 1.2.6 to 1.2.7 (:pr:`10666`)
+  - Fix ``test_categorize_info`` with nightly ``pyarrow`` (:pr:`10662`) `James Bourbeau`_
+
+  - Rewrite ``test_subprocess_cluster_does_not_depend_on_logging`` (:pr-distributed:`8409`) `Hendrik Makait`_
+  - Avoid ``RecursionError`` when failing to pickle key in ``SpillBuffer`` and using ``tblib=3`` (:pr-distributed:`8404`) `Hendrik Makait`_
+  - Allow tasks to override ``is_rootish`` heuristic  (:pr-distributed:`8412`) `Hendrik Makait`_
+  - Remove GPU executor (:pr-distributed:`8399`) `Hendrik Makait`_
+  - Do not rely on logging for subprocess cluster (:pr-distributed:`8398`) `Hendrik Makait`_
+  - Update gpuCI ``RAPIDS_VER`` to ``24.02`` (:pr-distributed:`8384`)
+  - Bump ``actions/setup-python`` from 4 to 5 (:pr-distributed:`8396`)
+  - Ensure output chunks in P2P rechunking are distributed homogeneously (:pr-distributed:`8207`) `Florian Jetter`_
+  - Trivial: fix typo (:pr-distributed:`8395`) `crusaderky`_
+  - Bump ``JamesIves/github-pages-deploy-action`` from 4.4.3 to 4.5.0 (:pr-distributed:`8387`)
+  - Bump ``conda-incubator/setup-miniconda from`` 3.0.0 to 3.0.1 (:pr-distributed:`8388`)
+
+
 .. _v2023.12.0:
 
 2023.12.0
@@ -7412,3 +7469,5 @@ Other
 .. _`Samantha Hughes`: https://github.com/shughes-uk
 .. _`Mario Šaško`: https://github.com/mariosasko
 .. _`joanrue`: https://github.com/joanrue
+.. _`Andrew S. Rosen`: https://github.com/Andrew-S-Rosen
+.. _`jochenott`: https://github.com/jochenott
diff --git a/pyproject.toml b/pyproject.toml
index 0ef22fb6d..cd2cfe9b7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -56,7 +56,7 @@ dataframe = [
     "dask[array]",
     "pandas >= 1.3",
 ]
-distributed = ["distributed == 2023.12.0"]
+distributed = ["distributed == 2023.12.1"]
 diagnostics = [
     "bokeh >= 2.4.2",
     "jinja2 >= 2.10.3",


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 0b518eee4..7f40fd027 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -14,7 +14,6 @@ np = pytest.importorskip("numpy")
 import math
 import operator
 import os
-import sys
 import time
 import warnings
 from functools import reduce
@@ -22,13 +21,11 @@ from io import StringIO
 from operator import add, sub
 from threading import Lock
 
-from packaging.version import Version
 from tlz import concat, merge
 from tlz.curried import identity
 
 import dask
 import dask.array as da
-from dask._compatibility import PY_VERSION
 from dask.array.chunk import getitem
 from dask.array.core import (
     Array,
@@ -3914,10 +3911,6 @@ def test_setitem_1d():
         dx[index] = 1
 
 
-@pytest.mark.xfail(
-    sys.platform == "win32" and PY_VERSION >= Version("3.12.0"),
-    reason="https://github.com/dask/dask/issues/10604",
-)
 def test_setitem_hardmask():
     x = np.ma.array([1, 2, 3, 4], dtype=int)
     x.harden_mask()
@@ -3931,7 +3924,20 @@ def test_setitem_hardmask():
     dx = da.from_array(y)
     dx[0] = np.ma.masked
     dx[0:2] = np.ma.masked
+    assert_eq(x, dx)
 
+
+def test_setitem_slice_twice():
+    x = np.array([1, 2, 3, 4, 5, 6], dtype=int)
+    val = np.array([0, 0], dtype=int)
+    y = x.copy()
+
+    x[0:2] = val
+    x[4:6] = val
+
+    dx = da.from_array(y)
+    dx[0:2] = val
+    dx[4:6] = val
     assert_eq(x, dx)
 
 
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 7c661f896..364e385eb 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2275,7 +2275,7 @@ def test_repartition_npartitions(use_index, n, k, dtype, transform):
     )
     df = transform(df)
     a = dd.from_pandas(df, npartitions=n, sort=use_index)
-    b = a.repartition(k)
+    b = a.repartition(npartitions=k)
     assert_eq(a, b)
     assert b.npartitions == k
     assert all(map(len, b.partitions))
@@ -2302,7 +2302,7 @@ def test_repartition_partition_size(use_index, n, partition_size, transform):
 def test_repartition_partition_size_arg():
     df = pd.DataFrame({"x": range(10)})
     a = dd.from_pandas(df, npartitions=2)
-    b = a.repartition("1 MiB")
+    b = a.repartition(partition_size="1 MiB")
     assert b.npartitions == 1
 
 
@@ -2356,7 +2356,7 @@ def test_repartition_datetime_tz_index():
     ds = dd.from_pandas(s, npartitions=2)
     assert ds.npartitions == 2
     assert_eq(s, ds)
-    result = ds.repartition(5)
+    result = ds.repartition(npartitions=5)
     assert result.npartitions == 5
     assert_eq(s, result)
 
diff --git a/dask/dataframe/tests/test_merge_column_and_index.py b/dask/dataframe/tests/test_merge_column_and_index.py
index b7f64b112..bd455b61d 100644
--- a/dask/dataframe/tests/test_merge_column_and_index.py
+++ b/dask/dataframe/tests/test_merge_column_and_index.py
@@ -240,7 +240,7 @@ def test_merge_column_with_nulls(repartition):
     df1_d = dd.from_pandas(df1, npartitions=4)
     df2_d = dd.from_pandas(df2, npartitions=3).set_index("b")
     if repartition:
-        df2_d = df2_d.repartition(repartition)
+        df2_d = df2_d.repartition(npartitions=repartition)
 
     pandas_result = df1.merge(
         df2.set_index("b"), how="left", left_on="a", right_index=True
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
index cdaa3a537..05bec5513 100644
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -1230,7 +1230,9 @@ def test_set_index_overlap_2():
         )
     )
     ddf = dd.from_pandas(df, npartitions=2)
-    result = ddf.reset_index().repartition(8).set_index("index", sorted=True)
+    result = (
+        ddf.reset_index().repartition(npartitions=8).set_index("index", sorted=True)
+    )
     expected = df.reset_index().set_index("index")
     assert_eq(result, expected)
     assert result.npartitions == 8
@@ -1281,7 +1283,7 @@ def test_compute_current_divisions_overlap_2():
         )
     )
     ddf1 = dd.from_pandas(data, npartitions=2)
-    ddf2 = ddf1.clear_divisions().repartition(8)
+    ddf2 = ddf1.clear_divisions().repartition(npartitions=8)
     with pytest.warns(UserWarning, match="Partitions have overlapping values"):
         ddf2.compute_current_divisions()