Repository: dask/dask
Pull request #10706 (2023-12-14 16:59:56)

[Dask.order] Ignore data tasks when ordering
This supersedes https://github.com/dask/dask/pull/10619

This may be a little controversial... However, there are frequently topologies (particularly in the array space) that have a dummy task at the bottom of the graph that includes some metadata (e.g. for zarr). In the xarray world, those are frequently embedded numpy arrays.

I believe we should special case such tasks since they can throw off otherwise fine heuristics.

1. They never have dependencies so we can schedule them whenever we want
2. There is no point in running their dependents more quickly than others trying to release them. We cannot release their data since the data is embedded in the graph/run_spec (even a released task currently holds on to their run_spec)
3. The data itself is typically very small, otherwise it would not be feasible to embed it into a graph

So, why is this controversial

With this, ordering would be different for say da.from_numpy(np.zeros(100), chunks=20) and da.zeros(100, chunk=20) since the first would literally embed the numpy array into the dask graph while the latter generates the data whenever needed. I'm not sure if this is such a bad thing. It may just be a little surprising but I don't think this will have negative effects. This can be seen in one of our tests where I changed to a `dask.array.random`

fjetter: I'm not aware of any cases that perform better or worse with this but I am reasonably certain that this graph transformation is a good thing, see the argument above. The only place I am aware of where this is used is really the array world where zarr et al have this dummy task.

This should put an end to the conversation about whether or not they should inline those tasks, at least for the dask/dask part. There is still the distributed scheduler that could mess this up, of course.
fjetter: Generally speaking, I could see us add various equivalency transformations to a graph that we know won't impact performance but will simplify the topology. This is now the second one I added, after https://github.com/dask/dask/pull/10697 which is basically an equivalent to this PR but for leaf nodes.
I am not aware of any such transformations at the moment but if more problems around ordering pop up, this might be an angle to investigate

Patch:
diff --git a/.github/labeler.yml b/.github/labeler.yml
deleted file mode 100644
index 5247cc239..000000000
--- a/.github/labeler.yml
+++ /dev/null
@@ -1,22 +0,0 @@
-dataframe:
-  - dask/dataframe/*
-  - dask/dataframe/**/*
-
-array:
-  - dask/array/*
-  - dask/array/**/*
-
-io:
-  - dask/dataframe/io/*
-  - dask/dataframe/io/**/*
-
-documentation:
-  - docs/*
-  - docs/**/*
-
-dispatch:
-  - dask/array/backends.py
-  - dask/array/dispatch.py
-  - dask/dataframe/backends.py
-  - dask/dataframe/dispatch.py
-  - dask/dataframe/extensions.py
diff --git a/.github/workflows/label-prs.yml b/.github/workflows/label-prs.yml
deleted file mode 100644
index 834bbcd32..000000000
--- a/.github/workflows/label-prs.yml
+++ /dev/null
@@ -1,12 +0,0 @@
-name: "PR Labeler"
-on:
-- pull_request_target
-
-jobs:
-  label:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/labeler@v5.0.0
-      with:
-        repo-token: "${{ secrets.GITHUB_TOKEN }}"
-        sync-labels: false
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/dask/dataframe/multi.py b/dask/dataframe/multi.py
index 82d43f315..a46530fd2 100644
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -1412,7 +1412,15 @@ def _split_partition(df, on, nsplits):
     # add a "_partitions" column to perform the split.
     if not isinstance(on, _Frame):
         on = _select_columns_or_index(df, on)
-    partitions = partitioning_index(on, nsplits)
+
+    dtypes = {}
+    for col, dtype in on.dtypes.items():
+        if pd.api.types.is_numeric_dtype(dtype):
+            dtypes[col] = np.float64
+    if not dtypes:
+        dtypes = None
+
+    partitions = partitioning_index(on, nsplits, cast_dtype=dtypes)
     df2 = df.assign(_partitions=partitions)
     return shuffle_group(
         df2,
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index d83e4d118..d99320828 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -420,6 +420,13 @@ def shuffle(
         # selection will not match (important when merging).
         index = index.to_frame()
 
+    dtypes = {}
+    for col, dtype in index.dtypes.items():
+        if pd.api.types.is_numeric_dtype(dtype):
+            dtypes[col] = np.float64
+    if not dtypes:
+        dtypes = None
+
     meta = df._meta._constructor_sliced([0])
     # Ensure that we have the same index as before to avoid alignment
     # when calculating meta dtypes later on
@@ -429,6 +436,7 @@ def shuffle(
         npartitions=npartitions or df.npartitions,
         meta=meta,
         transform_divisions=False,
+        cast_dtype=dtypes,
     )
     df2 = df.assign(_partitions=partitions)
     df2._meta.index.name = df._meta.index.name
@@ -786,7 +794,7 @@ def rearrange_by_column_tasks(
 ########################################################
 
 
-def partitioning_index(df, npartitions):
+def partitioning_index(df, npartitions, cast_dtype=None):
     """
     Computes a deterministic index mapping each record to a partition.
 
@@ -797,12 +805,18 @@ def partitioning_index(df, npartitions):
     df : DataFrame/Series/Index
     npartitions : int
         The number of partitions to group into.
+    cast_dtype : dtype, optional
+        The dtype to cast to to avoid nullability issues
 
     Returns
     -------
     partitions : ndarray
         An array of int64 values mapping each record to a partition.
     """
+    if cast_dtype is not None:
+        # Fixme: astype raises with strings in numeric columns, but raising
+        # here might be very noisy
+        df = df.astype(cast_dtype, errors="ignore")
     return hash_object_dispatch(df, index=False) % int(npartitions)
 
 
diff --git a/dask/delayed.py b/dask/delayed.py
index c599c5cb0..7b5439909 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -716,6 +716,10 @@ class DelayedLeaf(Delayed):
     def __doc__(self):
         return self._obj.__doc__
 
+    @property
+    def __wrapped__(self):
+        return self._obj
+
 
 class DelayedAttr(Delayed):
     __slots__ = ("_obj", "_attr")
diff --git a/dask/order.py b/dask/order.py
index da7d6082f..0592eeda2 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -46,6 +46,7 @@ graph types laid out very carefully to show the kinds of situations that often
 arise, and the order we would like to be determined.
 
 """
+import copy
 from collections import defaultdict, deque, namedtuple
 from collections.abc import Iterable, Mapping, MutableMapping
 from typing import Any, Callable, Literal, NamedTuple, overload
@@ -109,16 +110,16 @@ def order(
     expected_len = len(dsk)
 
     if dependencies is None:
+        dependencies_are_copy = True
         dependencies = {k: get_dependencies(dsk, k) for k in dsk}
+    else:
+        # Below we're removing items from the sets in this dict
+        # We need a deepcopy for that but we only want to do this if necessary
+        # since this only happens for special cases.
+        dependencies_are_copy = False
+        dependencies = dict(dependencies)
 
     dependents = reverse_dict(dependencies)
-    num_needed, total_dependencies = ndependencies(dependencies, dependents)
-    if len(total_dependencies) != len(dsk):
-        cycle = getcycle(dsk, None)
-        raise RuntimeError(
-            "Cycle detected between the following keys:\n  -> %s"
-            % "\n  -> ".join(str(x) for x in cycle)
-        )
 
     leaf_nodes = {k for k, v in dependents.items() if not v}
     root_nodes = {k for k, v in dependencies.items() if not v}
@@ -128,11 +129,20 @@ def order(
     # Normalize the graph by removing leaf nodes that are not actual tasks, see
     # for instance da.store where the task is merely an alias
     # to multiple keys, i.e. [key1, key2, ...,]
+
+    # Similarly, we are removing root nodes that are pure data tasks. Those task
+    # are embedded in the run_spec of a task and are not runnable. We have to
+    # assign a priority but their priority has no impact on performance.
+    # The removal of those tasks typically transforms the graph topology in a
+    # way that is simpler to handle
     all_tasks = False
     n_removed_leaves = 0
+    requires_data_task = defaultdict(set)
     while not all_tasks:
         all_tasks = True
         for leaf in list(leaf_nodes):
+            if leaf in root_nodes:
+                continue
             if (
                 not istask(dsk[leaf])
                 # Having a linear chain is fine
@@ -155,6 +165,30 @@ def order(
                     if not dependents[dep]:
                         leaf_nodes.add(dep)
 
+        for root in list(root_nodes):
+            if root in leaf_nodes:
+                continue
+            if not istask(dsk[root]) and len(dependents[root]) > 1:
+                if not dependencies_are_copy:
+                    dependencies_are_copy = True
+                    dependencies = copy.deepcopy(dependencies)
+                root_nodes.remove(root)
+                for dep in dependents[root]:
+                    requires_data_task[dep].add(root)
+                    dependencies[dep].remove(root)
+                    if not dependencies[dep]:
+                        root_nodes.add(dep)
+                del dsk[root]
+                del dependencies[root]
+                del dependents[root]
+
+    num_needed, total_dependencies = ndependencies(dependencies, dependents)
+    if len(total_dependencies) != len(dsk):
+        cycle = getcycle(dsk, None)
+        raise RuntimeError(
+            "Cycle detected between the following keys:\n  -> %s"
+            % "\n  -> ".join(str(x) for x in cycle)
+        )
     assert dependencies is not None
     roots_connected, max_dependents = _connecting_to_roots(dependencies, dependents)
     leafs_connected, _ = _connecting_to_roots(dependents, dependencies)
@@ -180,8 +214,8 @@ def order(
         except KeyError:
             assert dependencies is not None
             _sort_keys_cache[x] = rv = (
-                len(dependencies[x]),
                 total_dependencies[x],
+                len(dependencies[x]),
                 len(roots_connected[x]),
                 -max_dependents[x],
                 # Converting to str is actually faster than creating some
@@ -199,12 +233,15 @@ def order(
         nonlocal i
         while next_items:
             item = next_items.pop()
-            assert not num_needed[item]
             runnable_hull.discard(item)
             reachable_hull.discard(item)
             leaf_nodes.discard(item)
             if item in result:
                 continue
+
+            while requires_data_task[item]:
+                add_to_result(requires_data_task[item].pop())
+
             if return_stats:
                 result[item] = Order(i, crit_path_counter - _crit_path_counter_offset)
             else:
@@ -215,7 +252,7 @@ def order(
             # the final result since the `process_runnable` should produce
             # equivalent results regardless of the order in which runnable is
             # populated (not identical but equivalent)
-            for dep in dependents[item]:
+            for dep in dependents.get(item, ()):
                 num_needed[dep] -= 1
                 reachable_hull.add(dep)
                 if not num_needed[dep]:
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 1970ef812..2d726b472 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -11,6 +11,34 @@ Released on December 15, 2023
 Highlights
 ^^^^^^^^^^
 
+Logical Query Planning now available for Dask DataFrames
+""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+Dask DataFrames are now much more performant by using a logical query planner.
+This feature is currently off by default, but can be turned on with:
+
+.. code:: python
+
+    dask.config.set({"dataframe.query-planning": True})
+
+You also need to have ``dask-expr`` installed:
+
+.. code:: bash
+
+    pip install dask-expr
+
+We've seen promising performance improvements so far, see
+`this blog post <https://blog.coiled.io/blog/dask-expr-tpch-dask.html>`__
+and `these regularly updated benchmarks <https://tpch.coiled.io>`__  for more information.
+A more detailed explanation of how the query optimizer works can be found in
+`this blog post <https://blog.coiled.io/blog/dask-expr-introduction.html>`__.
+
+This feature is still under active development
+and the `API <https://github.com/dask-contrib/dask-expr#api-coverage>`__ isn't stable yet,
+so breaking changes can occur. We expect to make the query optimizer the default early next year.
+
+See :pr:`10634` from `Patrick Hoefler`_ for details.
+
 Dtype inference in ``read_parquet``
 """""""""""""""""""""""""""""""""""
 
@@ -20,6 +48,87 @@ original Arrow array, and thus avoid the conversion to NumPy object. Additionall
 ``read_parquet`` will no longer infer nested and binary types as strings, they will
 be stored in NumPy object arrays.
 
+See :pr:`10698` and :pr:`10705` from `Patrick Hoefler`_ for details.
+
+Scheduling improvements to reduce memory usage
+""""""""""""""""""""""""""""""""""""""""""""""
+
+This release includes a major rewrite to a core part of our scheduling logic. It
+includes a new approach to the topological sorting algorithm in ``dask.order``
+which determines the order in which tasks are run. Improper ordering is known to
+be a major contributor to too large cluster memory pressure.
+
+Updates in this release fix a couple of performance regressions that were introduced
+in the release ``2023.10.0`` (see :pr:`10535`). Generally, computations should now
+be much more eager to release data if it is no longer required in memory.
+
+See :pr:`10660`, :pr:`10697` from `Florian Jetter`_ for details.
+
+Improved P2P-based merging robustness and performance
+"""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+This release contains several updates that fix a possible deadlock introduced in
+2023.9.2 and improve the robustness of P2P-based merging when the cluster is
+dynamically scaling up.
+
+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:
 
@@ -37,6 +146,15 @@ PipInstall restart and environment variables
 The ``distributed.PipInstall`` plugin now has more robust restart logic and also supports
 `environment variables <https://pip.pypa.io/en/stable/reference/requirements-file-format/#using-environment-variables>`_.
 
+Below shows how users can use the ``distributed.PipInstall`` plugin and a ``TOKEN`` environment
+variable to securely install a package from a private repository:
+
+.. code:: python
+
+  from dask.distributed import PipInstall
+  plugin = PipInstall(packages=["private_package@git+https://${TOKEN}@github.com/dask/private_package.git])
+  client.register_plugin(plugin)
+
 See :pr-distributed:`8374`, :pr-distributed:`8357`, and :pr-distributed:`8343` from `Hendrik Makait`_ for details.
 
 
@@ -47,7 +165,7 @@ Previously the contents of dashboard plots wouldn't be displayed.
 
 See :pr-distributed:`8347` and :pr-distributed:`8381` from `Jacob Tomlinson`_ for details.
 
-  
+
 .. dropdown:: Additional changes
 
   - Add ``network`` marker to ``test_pyarrow_filesystem_option_real_data`` (:pr:`10653`) `Richard (Rick) Zamora`_
@@ -243,7 +361,7 @@ This release contains major updates to Dask's task graph scheduling logic.
 The updates here significantly reduce memory pressure on array reductions.
 We anticipate this will have a strong impact on the array computing community.
 
-See :pr-distributed:`10535` from `Florian Jetter`_ for details.
+See :pr:`10535` from `Florian Jetter`_ for details.
 
 
 Improved P2P shuffling robustness
@@ -7364,3 +7482,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/docs/source/dataframe-api.rst b/docs/source/dataframe-api.rst
index 73723fc2c..cbaa41076 100644
--- a/docs/source/dataframe-api.rst
+++ b/docs/source/dataframe-api.rst
@@ -1,5 +1,5 @@
-API
----
+Dask DataFrame API
+==================
 
 .. currentmodule:: dask.dataframe
 
diff --git a/docs/source/dataframe-create.rst b/docs/source/dataframe-create.rst
index 811953f75..c2cabf2cb 100644
--- a/docs/source/dataframe-create.rst
+++ b/docs/source/dataframe-create.rst
@@ -1,5 +1,5 @@
-Create and Store Dask DataFrames
-================================
+Load and Save Data with Dask DataFrames
+=======================================
 
 .. meta::
     :description: Learn how to create DataFrames and store them. Create a Dask DataFrame from various data storage formats like CSV, HDF, Apache Parquet, and others.
@@ -67,7 +67,7 @@ Read from CSV
 
 You can use :func:`read_csv` to read one or more CSV files into a Dask DataFrame.
 It supports loading multiple files at once using globstrings:
-  
+
 .. code-block:: python
 
    >>> df = dd.read_csv('myfiles.*.csv')
@@ -76,7 +76,7 @@ You can break up a single large file with the ``blocksize`` parameter:
 
 .. code-block:: python
 
-   >>> df = dd.read_csv('largefile.csv', blocksize=25e6)  # 25MB chunks  
+   >>> df = dd.read_csv('largefile.csv', blocksize=25e6)  # 25MB chunks
 
 Changing the ``blocksize`` parameter will change the number of partitions (see the explanation on
 :ref:`partitions <dataframe-design-partitions>`). A good rule of thumb when working with
diff --git a/docs/source/dataframe-design.rst b/docs/source/dataframe-design.rst
index 48a63bfbc..95046456e 100644
--- a/docs/source/dataframe-design.rst
+++ b/docs/source/dataframe-design.rst
@@ -1,7 +1,7 @@
 .. _dataframe.design:
 
-Internal Design
-===============
+Dask DataFrame Design
+=====================
 
 Dask DataFrames coordinate many Pandas DataFrames/Series arranged along an
 index.  We define a Dask DataFrame object with the following components:
diff --git a/docs/source/dataframe-extra.rst b/docs/source/dataframe-extra.rst
new file mode 100644
index 000000000..41126dd90
--- /dev/null
+++ b/docs/source/dataframe-extra.rst
@@ -0,0 +1,14 @@
+Additional Information
+======================
+
+.. toctree::
+   :maxdepth: 1
+
+   Parquet <dataframe-parquet.rst>
+   Indexing <dataframe-indexing.rst>
+   SQL <dataframe-sql.rst>
+   Join Performance <dataframe-joins.rst>
+   Shuffling Performance <dataframe-groupby.rst>
+   dataframe-categoricals.rst
+   Extend <dataframe-extend.rst>
+   Hive Partitioning <dataframe-hive.rst>
diff --git a/docs/source/dataframe-groupby.rst b/docs/source/dataframe-groupby.rst
index 5b81f90ac..03261b3ad 100644
--- a/docs/source/dataframe-groupby.rst
+++ b/docs/source/dataframe-groupby.rst
@@ -1,5 +1,5 @@
-Shuffling for GroupBy and Join
-==============================
+Shuffling Performance
+=====================
 
 .. currentmodule:: dask.dataframe
 
diff --git a/docs/source/dataframe.rst b/docs/source/dataframe.rst
index 1278411c6..a80dc9d59 100644
--- a/docs/source/dataframe.rst
+++ b/docs/source/dataframe.rst
@@ -8,198 +8,181 @@ Dask DataFrame
    :maxdepth: 1
    :hidden:
 
-   dataframe-create.rst
+   Load and Save Data <dataframe-create.rst>
+   Internal Design <dataframe-design.rst>
    Best Practices <dataframe-best-practices.rst>
-   dataframe-design.rst
-   dataframe-groupby.rst
-   dataframe-joins.rst
-   dataframe-indexing.rst
-   dataframe-categoricals.rst
-   dataframe-extend.rst
-   dataframe-parquet.rst
-   dataframe-hive.rst
-   dataframe-sql.rst
-   dataframe-api.rst
-
-A Dask DataFrame is a large parallel DataFrame composed of many smaller pandas
-DataFrames, split along the index.  These pandas DataFrames may live on disk
-for larger-than-memory computing on a single machine, or on many different
-machines in a cluster.  One Dask DataFrame operation triggers many operations
-on the constituent pandas DataFrames.
+   API <dataframe-api.rst>
+   dataframe-extra.rst
 
-.. raw:: html
+.. grid:: 1 1 2 2
 
-   <iframe width="560"
-           height="315"
-           src="https://www.youtube.com/embed/AT2XtFehFSQ"
-           style="margin: 0 auto 20px auto; display: block;"
-           frameborder="0"
-           allow="accelerometer; autoplay; encrypted-media; gyroscope; picture-in-picture"
-           allowfullscreen></iframe>
+   .. grid-item::
+      :columns: 12 12 8 8
 
-Examples
---------
+      Dask DataFrame helps you process large tabular data by parallelizing pandas,
+      either on your laptop for larger-than-memory computing, or on a distributed
+      cluster of computers.
 
-Visit https://examples.dask.org/dataframe.html to see and run examples using
-Dask DataFrame.
+      -  **Just pandas:** Dask DataFrames are a collection of many pandas DataFrames.
 
+         The API is the same.  The execution is the same.
+      -  **Large scale:** Works on 100 GiB on a laptop, or 100 TiB on a cluster.
+      -  **Easy to use:** Pure Python, easy to set up and debug.
 
-.. image:: images/dask-dataframe.svg
-   :alt: Column of four squares collectively labeled as a Dask DataFrame with a single constituent square labeled as a pandas DataFrame.
-   :width: 45%
-   :align: right
+   .. grid-item::
+      :columns: 12 12 4 4
 
-Design
-------
+      .. image:: images/dask-dataframe.svg
+         :alt: Column of four squares collectively labeled as a Dask DataFrame with a single constituent square labeled as a pandas DataFrame.
 
 Dask DataFrames coordinate many pandas DataFrames/Series arranged along the
 index.  A Dask DataFrame is partitioned *row-wise*, grouping rows by index value
 for efficiency.  These pandas objects may live on disk or on other machines.
 
-Dask DataFrame copies the pandas DataFrame API
-----------------------------------------------
-
-Because the ``dask.DataFrame`` application programming interface (API) is a
-subset of the ``pd.DataFrame`` API, it should be familiar to pandas users.
-There are some slight alterations due to the parallel nature of Dask:
-
-.. grid:: 2
-
-    .. grid-item-card:: Dask DataFrame API
-
-        .. code-block:: python
-
-            >>> import dask.dataframe as dd
-            >>> df = dd.read_csv('2014-*.csv')
-            >>> df.head()
-                x  y
-            0  1  a
-            1  2  b
-            2  3  c
-            3  4  a
-            4  5  b
-            5  6  c
-
-            >>> df2 = df[df.y == 'a'].x + 1
-            >>> df2.compute()
-            0    2
-            3    5
-            Name: x, dtype: int64
-
-    .. grid-item-card:: pandas DataFrame API
-
-        .. code-block:: python
-
-            >>> import pandas as pd
-            >>> df = pd.read_csv('2014-1.csv')
-            >>> df.head()
-                x  y
-            0  1  a
-            1  2  b
-            2  3  c
-            3  4  a
-            4  5  b
-            5  6  c
-
-            >>> df2 = df[df.y == 'a'].x + 1
-            >>> df2
-            0    2
-            3    5
-            Name: x, dtype: int64
+From pandas to Dask
+-------------------
+
+Dask DataFrame copies pandas, and so should be familiar to most users
+
+.. tab-set::
+
+   .. tab-item:: Load Data
+
+      Pandas and Dask have the same API, and so switching from one to the other
+      is straightforward.
+
+      .. grid:: 1 1 2 2
+
+          .. grid-item::
+
+              .. code-block:: python
+
+                  >>> import pandas as pd
+
+                  >>> df = pd.read_parquet('s3://mybucket/myfile.parquet')
+                  >>> df.head()
+                  0  1  a
+                  1  2  b
+                  2  3  c
+
+          .. grid-item::
+
+              .. code-block:: python
+
+                  >>> import dask.dataframe as dd
+
+                  >>> df = dd.read_parquet('s3://mybucket/myfile.*.parquet')
+                  >>> df.head()
+                  0  1  a
+                  1  2  b
+                  2  3  c
+
+   .. tab-item:: Data Processing
+
+      Dask does pandas in parallel.
+      Dask is lazy; when you want an in-memory result add ``.compute()``.
+
+      .. grid:: 1 1 2 2
+
+          .. grid-item::
+
+              .. code-block:: python
+
+                  >>> import pandas as pd
+
+                  >>> df = df[df.value >= 0]
+                  >>> joined = df.merge(other, on="account")
+                  >>> result = joined.groupby("account")
+
+                  >>> result
+                  alice 123
+                  bob   456
+
+
+          .. grid-item::
+
+              .. code-block:: python
+
+                  >>> import dask.dataframe as dd
+
+                  >>> df = df[df.value >= 0]
+                  >>> joined = df.merge(other, on="account")
+                  >>> result = joined.groupby("account").value.mean()
+
+                  >>> result.compute()
+                  alice 123
+                  bob   456
+
+   .. tab-item:: Machine Learning
+
+      Machine learning libraries often have Dask submodules that
+      expect Dask DataFrames and operate in parallel.
+
+      .. grid:: 1 1 2 2
+
+          .. grid-item::
+
+              .. code-block:: python
+
+                  >>> import pandas as pd
+                  >>> import xgboost
+                  >>> from sklearn.cross_validation import train_test_split
+
+                  >>> X_train, X_test, y_train, y_test = train_test_split(
+                  ...    X, y, test_size=0.2,
+                  )
+                  >>> dtrain = xgboost.DMatrix(X_train, label=y_train)
+
+                  >>> xgboost.train(params, dtrain, 100)
+                  <xgboost.Booster ...>
+
+          .. grid-item::
+
+              .. code-block:: python
+
+                  >>> import dask.dataframe as dd
+                  >>> import xgboost.dask
+                  >>> from dask_ml.model_selection import train_test_split
+
+                  >>> X_train, X_test, y_train, y_test = train_test_split(
+                  ...    X, y, test_size=0.2,
+                  )
+                  >>> dtrain = xgboost.dask.DaskDMatrix(client, X, y)
+
+                  >>> xgboost.dask.train(params, dtrain, 100)
+                  <xgboost.Booster ...>
 
 As with all Dask collections, you trigger computation by calling the
-``.compute()`` method.
-
-Common Uses and Anti-Uses
--------------------------
-
-Dask DataFrame is used in situations where pandas is commonly needed, usually when
-pandas fails due to data size or computation speed:
-
--  Manipulating large datasets, even when those datasets don't fit in memory
--  Accelerating long computations by using many cores
--  Distributed computing on large datasets with standard pandas operations like
-   groupby, join, and time series computations
-
-Dask DataFrame may not be the best choice in the following situations:
-
-*  If your dataset fits comfortably into RAM on your laptop, then you may be
-   better off just using pandas.  There may be simpler ways to improve
-   performance than through parallelism
-*  If your dataset doesn't fit neatly into the pandas tabular model, then you
-   might find more use in :doc:`dask.bag <bag>` or :doc:`dask.array <array>`
-*  If you need functions that are not implemented in Dask DataFrame, then you
-   might want to look at :doc:`dask.delayed <delayed>` which offers more
-   flexibility
-*  If you need a proper database with all that databases offer you might prefer
-   something like Postgres_
-
-.. _Postgres: https://www.postgresql.org/
-
-
-Scope
------
-
-Dask DataFrame covers a well-used portion of the pandas API.
-The following class of computations works well:
-
-* Trivially parallelizable operations (fast):
-    *  Element-wise operations:  ``df.x + df.y``, ``df * df``
-    *  Row-wise selections:  ``df[df.x > 0]``
-    *  Loc:  ``df.loc[4.0:10.5]``
-    *  Common aggregations:  ``df.x.max()``, ``df.max()``
-    *  Is in:  ``df[df.x.isin([1, 2, 3])]``
-    *  Date time/string accessors:  ``df.timestamp.month``
-* Cleverly parallelizable operations (fast):
-    *  groupby-aggregate (with common aggregations): ``df.groupby(df.x).y.max()``,
-       ``df.groupby('x').min()`` (see :ref:`dataframe.groupby.aggregate`)
-    *  groupby-apply on index: ``df.groupby(['idx', 'x']).apply(myfunc)``, where
-       ``idx`` is the index level name
-    *  value_counts:  ``df.x.value_counts()``
-    *  Drop duplicates:  ``df.x.drop_duplicates()``
-    *  Join on index:  ``dd.merge(df1, df2, left_index=True, right_index=True)``
-       or ``dd.merge(df1, df2, on=['idx', 'x'])`` where ``idx`` is the index
-       name for both ``df1`` and ``df2``
-    *  Join with pandas DataFrames: ``dd.merge(df1, df2, on='id')``
-    *  Element-wise operations with different partitions / divisions: ``df1.x + df2.y``
-    *  Date time resampling: ``df.resample(...)``
-    *  Rolling averages:  ``df.rolling(...)``
-    *  Pearson's correlation: ``df[['col1', 'col2']].corr()``
-* Operations requiring a shuffle (slow-ish, unless on index, see :doc:`dataframe-groupby`)
-    *  Set index:  ``df.set_index(df.x)``
-    *  groupby-apply not on index (with anything):  ``df.groupby(df.x).apply(myfunc)``
-    *  Join not on the index:  ``dd.merge(df1, df2, on='name')``
-
-However, Dask DataFrame does not implement the entire pandas interface.  Users
-expecting this will be disappointed.  Notably, Dask DataFrame has the following
-limitations:
-
-1.  Setting a new index from an unsorted column is expensive
-2.  Many operations like groupby-apply and join on unsorted columns require
-    setting the index, which as mentioned above, is expensive
-3.  The pandas API is very large.  Dask DataFrame does not attempt to implement
-    many pandas features or any of the more exotic data structures like NDFrames
-4.  Operations that were slow on pandas, like iterating through row-by-row,
-    remain slow on Dask DataFrame
-
-See the :doc:`DataFrame API documentation<dataframe-api>` for a more extensive list.
-
-
-Execution
----------
-
-By default, Dask DataFrame uses the :ref:`multi-threaded scheduler <threaded-scheduler>`.
-This exposes some parallelism when pandas or the underlying NumPy operations
-release the global interpreter lock (GIL).  Generally, pandas is more GIL
-bound than NumPy, so multi-core speed-ups are not as pronounced for
-Dask DataFrame as they are for Dask Array. This is particularly true
-for string-heavy Python DataFrames, as Python strings are GIL bound.
-
-There has been recent work on changing the underlying representation
-of pandas string data types to be backed by
-`PyArrow Buffers <https://arrow.apache.org/docs/cpp/memory.html#buffers>`_, which
-should release the GIL, however, this work is still considered experimental.
-
-When dealing with text data, you may see speedups by switching to the
-:doc:`distributed scheduler <deploying>` either on a cluster or
-single machine.
+``.compute()`` method or persist data in distributed memory with the
+``.persist()`` method.
+
+When not to use Dask DataFrames
+-------------------------------
+
+Dask DataFrames are often used either when ...
+
+1.  Your data is too big
+2.  Your computation is too slow and other techniques don't work
+
+You should probably stick to just using pandas if ...
+
+1.  Your data is small
+2.  Your computation is fast (subsecond)
+3.  There are simpler ways to accelerate your computation, like avoiding
+    ``.apply`` or Python for loops and using a built-in pandas method instead.
+
+Examples
+--------
+
+Visit https://examples.dask.org/dataframe.html to see and run examples using
+Dask DataFrame.
+
+.. raw:: html
+
+  <iframe width="560"
+          height="315"
+          src="https://www.youtube.com/embed/AT2XtFehFSQ"
+          style="margin: 0 auto 20px auto; display: block;"
+          frameborder="0"
+          allow="accelerometer; autoplay; encrypted-media; gyroscope; picture-in-picture"
+          allowfullscreen></iframe>
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 4f08b063d..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()
 
@@ -1693,3 +1695,25 @@ def test_set_index_pyarrow_dtype(data, dtype):
     pdf_result = pdf.set_index("arrow_col")
     ddf_result = ddf.set_index("arrow_col")
     assert_eq(ddf_result, pdf_result)
+
+
+def test_shuffle_nulls_introduced():
+    df1 = pd.DataFrame([[True], [False]] * 50, columns=["A"])
+    df1["B"] = list(range(100))
+
+    df2 = pd.DataFrame(
+        [[2, 3], [109, 2], [345, 3], [50, 7], [95, 1]], columns=["B", "C"]
+    )
+
+    ddf1 = dd.from_pandas(df1, npartitions=10)
+    ddf2 = dd.from_pandas(df2, npartitions=1)
+    meta = pd.Series(dtype=int, index=pd.Index([], dtype=bool, name="A"), name="A")
+    result = (
+        dd.merge(ddf1, ddf2, how="outer", on="B")
+        .groupby("A")
+        .apply(lambda df: len(df), meta=meta)
+    )
+    expected = (
+        pd.merge(df1, df2, how="outer", on="B").groupby("A").apply(lambda df: len(df))
+    )
+    assert_eq(result, expected, check_names=False)
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index fbb8dc39f..dc9b796ad 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -855,3 +855,4 @@ def test_delayed_function_attributes_forwarded():
 
     assert add.__name__ == "add"
     assert add.__doc__ == "This is a docstring"
+    assert add.__wrapped__(1, 2) == 3
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index c25e9e769..9b528c562 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -84,6 +84,7 @@ def test_ordering_keeps_groups_together(abcde):
     d = {(a, i): (f,) for i in range(4)}
     d.update({(b, 0): (f, (a, 0), (a, 1)), (b, 1): (f, (a, 2), (a, 3))})
     o = order(d)
+    assert_topological_sort(d, o)
 
     assert abs(o[(a, 0)] - o[(a, 1)]) == 1
     assert abs(o[(a, 2)] - o[(a, 3)]) == 1
@@ -91,6 +92,7 @@ def test_ordering_keeps_groups_together(abcde):
     d = {(a, i): (f,) for i in range(4)}
     d.update({(b, 0): (f, (a, 0), (a, 2)), (b, 1): (f, (a, 1), (a, 3))})
     o = order(d)
+    assert_topological_sort(d, o)
 
     assert abs(o[(a, 0)] - o[(a, 2)]) == 1
     assert abs(o[(a, 1)] - o[(a, 3)]) == 1
@@ -116,6 +118,7 @@ def test_avoid_broker_nodes(abcde):
         (b, 2): (f, (a, 1)),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o[(a, 1)] < o[(b, 0)] or (o[(b, 1)] < o[(a, 0)] and o[(b, 2)] < o[(a, 0)])
 
 
@@ -142,6 +145,7 @@ def test_base_of_reduce_preferred(abcde):
     dsk[c] = 1
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     assert o[(b, 0)] <= 1
     assert o[(b, 1)] <= 3
@@ -175,6 +179,7 @@ def test_avoid_upwards_branching(abcde):
     }
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     assert o[(d, 1)] < o[(b, 1)]
 
@@ -214,6 +219,7 @@ def test_avoid_upwards_branching_complex(abcde):
     }
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o[(c, 1)] < o[(b, 1)]
     assert abs(o[(d, 2)] - o[(d, 3)]) == 1
 
@@ -238,6 +244,7 @@ def test_deep_bases_win_over_dependents(abcde):
     dsk = {a: (f, b, c, d), b: (f, d, e), c: (f, d), d: 1, e: 2}
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o[b] < o[c]
 
 
@@ -255,6 +262,7 @@ def test_prefer_deep(abcde):
     dsk = {a: 1, b: (f, a), c: (f, b), d: 1, e: (f, d)}
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o[a] < o[d]
     assert o[b] < o[d]
 
@@ -272,6 +280,7 @@ def test_break_ties_by_str(abcde):
     dsk["y"] = (f, list(x_keys))
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     expected = {"y": 10}
     expected.update({k: i for i, k in enumerate(x_keys[::-1])})
 
@@ -293,6 +302,7 @@ def test_gh_3055():
 
     dsk = dict(w.__dask_graph__())
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert max(diagnostics(dsk, o=o)[1]) <= 8
     L = [o[k] for k in w.__dask_keys__()]
     assert sum(x < len(o) / 2 for x in L) > len(L) / 3  # some complete quickly
@@ -323,6 +333,7 @@ def test_favor_longest_critical_path(abcde):
     dsk = {c: (f,), d: (f, c), e: (f, c), b: (f, c), a: (f, b)}
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o[d] > o[b]
     assert o[e] > o[b]
 
@@ -351,6 +362,7 @@ def test_run_smaller_sections(abcde):
         dd: (f, cc),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert max(diagnostics(dsk)[1]) <= 4  # optimum is 3
     # This is a mildly ambiguous example
     # https://github.com/dask/dask/pull/10535/files#r1337528255
@@ -399,7 +411,8 @@ def test_local_parents_of_reduction(abcde):
         c1: (f(c1), c2),
     }
 
-    order(dsk)
+    o = order(dsk)
+    assert_topological_sort(dsk, o)
     dask.get(dsk, [a1, b1, c1])  # trigger computation
 
     assert log == expected
@@ -436,6 +449,7 @@ def test_nearest_neighbor(abcde):
     }
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     assert 3 < sum(o[a + i] < len(o) / 2 for i in "123456789") < 7
     assert 1 < sum(o[b + i] < len(o) / 2 for i in "1234") < 4
@@ -445,6 +459,7 @@ def test_string_ordering():
     """Prefer ordering tasks by name first"""
     dsk = {("a", 1): (f,), ("a", 2): (f,), ("a", 3): (f,)}
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o == {("a", 1): 0, ("a", 2): 1, ("a", 3): 2} or o == {
         ("a", 1): 2,
         ("a", 2): 1,
@@ -462,6 +477,7 @@ def test_string_ordering_dependents():
     # See comment in add_to_result
     dsk = {("a", 1): (f, "b"), ("a", 2): (f, "b"), ("a", 3): (f, "b"), "b": (f,)}
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o == {"b": 0, ("a", 1): 1, ("a", 2): 2, ("a", 3): 3} or o == {
         "b": 0,
         ("a", 1): 3,
@@ -483,6 +499,7 @@ def test_prefer_short_narrow(abcde):
         (c, 2): (f, (c, 1), (a, 1), (b, 1)),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o[(b, 0)] < o[(b, 1)]
     assert o[(b, 0)] < o[(c, 2)]
     assert o[(c, 1)] < o[(c, 2)]
@@ -540,6 +557,7 @@ def test_prefer_short_ancestor(abcde):
         (c, 2): (f, (c, 1), (a, 1), (b, 1)),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     assert o[(a, 0)] < o[(a, 1)]
     assert o[(b, 0)] < o[(b, 1)]
@@ -581,6 +599,7 @@ def test_map_overlap(abcde):
     }
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     assert o[(b, 1)] < o[(e, 5)] or o[(b, 5)] < o[(e, 1)]
 
@@ -614,6 +633,7 @@ def test_use_structure_not_keys(abcde):
         (b, 0): (f, (a, 3), (a, 8), (a, 1)),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     assert max(diagnostics(dsk, o=o)[1]) == 3
     As = sorted(val for (letter, _), val in o.items() if letter == a)
@@ -642,6 +662,7 @@ def test_dont_run_all_dependents_too_early(abcde):
         dsk[(c, i)] = (f, (c, 0))
         dsk[(d, i)] = (f, (d, i - 1), (b, i), (c, i))
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     expected = [3, 6, 9, 12, 15, 18, 21, 24, 27, 30]
     actual = sorted(v for (letter, num), v in o.items() if letter == d)
@@ -680,6 +701,7 @@ def test_many_branches_use_ndependencies(abcde):
         (a, 3): (f, (a, 2), (b, 2), (c, 1), (dd, 3), (ee, 3)),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     # run all d's and e's first
     expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
     actual = sorted(v for (letter, _), v in o.items() if letter in {d, dd, e, ee})
@@ -742,6 +764,7 @@ def test_order_with_equal_dependents(abcde):
                 }
             )
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     total = 0
     for x in abc:
         for i in range(len(abc)):
@@ -757,6 +780,7 @@ def test_order_with_equal_dependents(abcde):
         for i in range(len(abc)):
             dsk2[(x, 7, i, 0)] = (f, (x, 6, i, 0))
     o = order(dsk2)
+    assert_topological_sort(dsk2, o)
     total = 0
     for x in abc:
         for i in range(len(abc)):
@@ -772,6 +796,7 @@ def test_order_with_equal_dependents(abcde):
         for i in range(len(abc)):
             del dsk3[(x, 6, i, 1)]
     o = order(dsk3)
+    assert_topological_sort(dsk3, o)
     total = 0
     for x in abc:
         for i in range(len(abc)):
@@ -787,6 +812,7 @@ def test_order_with_equal_dependents(abcde):
         for i in range(len(abc)):
             del dsk4[(x, 6, i, 0)]
     o = order(dsk4)
+    assert_topological_sort(dsk4, o)
     pressure = diagnostics(dsk4, o=o)[1]
     assert max(pressure) <= max_pressure
     for x in abc:
@@ -848,6 +874,7 @@ def test_terminal_node_backtrack():
         ),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert o[("a", 2)] < o[("a", 3)]
 
 
@@ -857,9 +884,8 @@ def test_array_store_final_order(tmpdir):
     # but with the graph actually generated by da.store.
     da = pytest.importorskip("dask.array")
     zarr = pytest.importorskip("zarr")
-    import numpy as np
 
-    arrays = [da.from_array(np.ones((110, 4)), chunks=(100, 2)) for i in range(4)]
+    arrays = [da.ones((110, 4), chunks=(100, 2)) for i in range(4)]
     x = da.concatenate(arrays, axis=0).rechunk((100, 2))
 
     store = zarr.DirectoryStore(tmpdir)
@@ -867,6 +893,7 @@ def test_array_store_final_order(tmpdir):
     dest = root.empty_like(name="dest", data=x, chunks=x.chunksize, overwrite=True)
     d = x.store(dest, lock=False, compute=False)
     o = order(d.dask)
+    assert_topological_sort(dict(d.dask), o)
     # Find the lowest store. Dask starts here.
     stores = [k for k in o if isinstance(k, tuple) and k[0].startswith("store-map-")]
     first_store = min(stores, key=lambda k: o[k])
@@ -968,8 +995,8 @@ def test_eager_to_compute_dependent_to_free_parent():
         "a70": (f, f),
         "a71": (f, f),
     }
-    dependencies, dependents = get_deps(dsk)
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     _, pressure = diagnostics(dsk, o=o)
     assert max(pressure) <= 8
@@ -1015,6 +1042,7 @@ def test_diagnostics(abcde):
         (e, 1): (f, (e, 0)),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     info, memory_over_time = diagnostics(dsk)
     # this is ambiguous, depending on whether we start from left or right
     assert all(o[key] == val.order for key, val in info.items())
@@ -1150,6 +1178,7 @@ def test_xarray_like_reduction():
             }
         )
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     _, pressure = diagnostics(dsk, o=o)
     assert max(pressure) <= 9
 
@@ -1298,6 +1327,7 @@ def test_anom_mean_raw(abcde):
     }
 
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     # The left hand computation branch should complete before we start loading
     # more data
     nodes_to_finish_before_loading_more_data = [
@@ -1751,6 +1781,7 @@ def test_flox_reduction(abcde):
         (g, 2, 2): (f, (a, 0), (e, 1)),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     of1 = list(o[(g, 1, ix)] for ix in range(3))
     of2 = list(o[(g, 2, ix)] for ix in range(3))
     assert max(of1) < min(of2) or max(of2) < min(of1)
@@ -1788,7 +1819,12 @@ def test_reduce_with_many_common_dependents(optimize, keep_self, ndeps, n_reduce
         dsk = collections_to_dsk([graph], optimize_graph=optimize)
     dependencies, dependents = get_deps(dsk)
     # Verify assumptions
+    before = len(dsk)
+    before_dsk = dsk.copy()
     o = order(dsk)
+    assert_topological_sort(dsk, o)
+    assert before == len(o) == len(dsk)
+    assert before_dsk == dsk
     # Verify assumptions (specifically that the reducers are sum-aggregate)
     assert ({"object", "sum", "sum-aggregate"}).issubset({key_split(k) for k in o})
     reducers = {k for k in o if key_split(k) == "sum-aggregate"}
@@ -1806,6 +1842,15 @@ def test_reduce_with_many_common_dependents(optimize, keep_self, ndeps, n_reduce
             )
 
 
+def assert_topological_sort(dsk, order):
+    dependencies, dependents = get_deps(dsk)
+    num_needed = {k: len(dependencies[k]) for k in dsk}
+    for k in sorted(dsk, key=order.__getitem__):
+        assert num_needed[k] == 0
+        for dep in dependents[k]:
+            num_needed[dep] -= 1
+
+
 def test_doublediff(abcde):
     a, b, c, d, e = abcde
     dsk = {
@@ -1925,6 +1970,7 @@ def test_gh_3055_explicit(abcde):
     assert len(con_r) == len(dsk)
     assert con_r[(e, 0)] == {("root", 0), (a, 1)}
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     assert max(diagnostics(dsk, o=o)[1]) <= 5
     assert o[(e, 0)] < o[(a, 3)] < o[(a, 4)]
     assert o[(a, 2)] < o[(a, 3)] < o[(a, 4)]
@@ -1959,6 +2005,7 @@ def test_order_flox_reduction_2(abcde):
         (d, 1, 1): (c, 1, 1),
     }
     o = order(dsk)
+    assert_topological_sort(dsk, o)
     final_nodes = sorted(
         [(d, ix, jx) for ix in range(2) for jx in range(2)], key=o.__getitem__
     )
@@ -2004,6 +2051,7 @@ def test_xarray_8414():
     assert len(shared_roots) == 1
     shared_root = shared_roots.pop()
     o = order(dsk)
+    assert_topological_sort(dsk, o)
 
     previous = None
     step = 0
@@ -2137,8 +2185,12 @@ def test_do_not_mutate_input():
         "b": (f, 1),
         "c": np.array([[1, 2], [3, 4]]),
         "d": ["a", "b", "c"],
+        "e": (f, "d"),
     }
+    dependencies, __build_class__ = get_deps(dsk)
+    dependencies_copy = dependencies.copy()
     dsk_copy = dsk.copy()
-    o = order(dsk)
-    assert o["d"] == len(dsk) - 1
-    assert len(dsk) == len(dsk_copy)
+    o = order(dsk, dependencies=dependencies)
+    assert_topological_sort(dsk, o)
+    assert dsk == dsk_copy
+    assert dependencies == dependencies_copy