Repository: dask/dask
Pull request #10552 (2023-10-10 00:20:39)

Dispatch the ``partd.Encode`` class used for disk-based shuffling
This change allows dask-cudf to register the necessary dispatch function (`partd_encode_dispatch`) for cudf-backed data to work with `shuffle="disk"`. Right now, the disk-based shuffle is hardcoded to use `partd.PandasBlocks` (a subclass of `partd.Encode`).


Patch:
diff --git a/dask/core.py b/dask/core.py
index 8c5908751..ef75ca563 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -1,10 +1,10 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from collections.abc import Collection, Iterable
-from typing import Any, cast
+from collections.abc import Collection, Iterable, Mapping
+from typing import Any, Literal, TypeVar, cast, overload
 
-from dask.typing import Key, no_default
+from dask.typing import Graph, Key, NoDefault, no_default
 
 
 def ishashable(x):
@@ -223,7 +223,32 @@ def validate_key(key: object) -> None:
         raise TypeError(f"Unexpected key type {type(key)} (value: {key!r})")
 
 
-def get_dependencies(dsk, key=None, task=no_default, as_list=False):
+@overload
+def get_dependencies(
+    dsk: Graph,
+    key: Key | None = ...,
+    task: Key | NoDefault = ...,
+    as_list: Literal[False] = ...,
+) -> set[Key]:
+    ...
+
+
+@overload
+def get_dependencies(
+    dsk: Graph,
+    key: Key | None,
+    task: Key | NoDefault,
+    as_list: Literal[True],
+) -> list[Key]:
+    ...
+
+
+def get_dependencies(
+    dsk: Graph,
+    key: Key | None = None,
+    task: Key | NoDefault = no_default,
+    as_list: bool = False,
+) -> set[Key] | list[Key]:
     """Get the immediate tasks on which this task depends
 
     Examples
@@ -264,7 +289,7 @@ def get_dependencies(dsk, key=None, task=no_default, as_list=False):
     return keys_in_tasks(dsk, [arg], as_list=as_list)
 
 
-def get_deps(dsk):
+def get_deps(dsk: Graph) -> tuple[dict[Key, set[Key]], dict[Key, set[Key]]]:
     """Get dependencies and dependents from dask dask graph
 
     >>> inc = lambda x: x + 1
@@ -308,7 +333,10 @@ def flatten(seq, container=list):
                 yield item
 
 
-def reverse_dict(d):
+T_ = TypeVar("T_")
+
+
+def reverse_dict(d: Mapping[T_, Iterable[T_]]) -> dict[T_, set[T_]]:
     """
 
     >>> a, b, c = 'abc'
@@ -316,14 +344,13 @@ def reverse_dict(d):
     >>> reverse_dict(d)  # doctest: +SKIP
     {'a': set([]), 'b': set(['a']}, 'c': set(['a', 'b'])}
     """
-    result = defaultdict(set)
+    result: defaultdict[T_, set[T_]] = defaultdict(set)
     _add = set.add
     for k, vals in d.items():
         result[k]
         for val in vals:
             _add(result[val], k)
-    result.default_factory = None
-    return result
+    return dict(result)
 
 
 def subs(task, key, val):
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
index 9ef5d1b36..ee900c3bf 100644
--- a/dask/dataframe/backends.py
+++ b/dask/dataframe/backends.py
@@ -27,6 +27,7 @@ from dask.dataframe.dispatch import (
     make_meta_obj,
     meta_lib_from_array,
     meta_nonempty,
+    partd_encode_dispatch,
     pyarrow_schema_dispatch,
     to_pandas_dispatch,
     to_pyarrow_table_dispatch,
@@ -242,6 +243,13 @@ def get_pandas_dataframe_from_pyarrow(meta, table, **kwargs):
     return table.to_pandas(types_mapper=types_mapper, **kwargs)
 
 
+@partd_encode_dispatch.register(pd.DataFrame)
+def partd_pandas_blocks(_):
+    from partd import PandasBlocks
+
+    return PandasBlocks
+
+
 @meta_nonempty.register(pd.DatetimeTZDtype)
 @make_meta_dispatch.register(pd.DatetimeTZDtype)
 def make_meta_pandas_datetime_tz(x, index=None):
diff --git a/dask/dataframe/dispatch.py b/dask/dataframe/dispatch.py
index 307d5269a..9c70302de 100644
--- a/dask/dataframe/dispatch.py
+++ b/dask/dataframe/dispatch.py
@@ -25,6 +25,7 @@ tolist_dispatch = Dispatch("tolist")
 is_categorical_dtype_dispatch = Dispatch("is_categorical_dtype")
 union_categoricals_dispatch = Dispatch("union_categoricals")
 grouper_dispatch = Dispatch("grouper")
+partd_encode_dispatch = Dispatch("partd_encode_dispatch")
 pyarrow_schema_dispatch = Dispatch("pyarrow_schema_dispatch")
 from_pyarrow_table_dispatch = Dispatch("from_pyarrow_table_dispatch")
 to_pyarrow_table_dispatch = Dispatch("to_pyarrow_table_dispatch")
diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
index f00accde9..584279e73 100644
--- a/dask/dataframe/io/parquet/core.py
+++ b/dask/dataframe/io/parquet/core.py
@@ -317,10 +317,8 @@ def read_parquet(
         set to 'infer' or 'adaptive'. Default may be engine-dependant, but is
         128 MiB for the 'pyarrow' and 'fastparquet' engines.
     aggregate_files : bool or str, default None
-        WARNING: The ``aggregate_files`` argument will be deprecated in the future.
-        Please consider using ``from_map`` to create a DataFrame collection with a
-        custom file-to-partition mapping. If you strongly oppose the deprecation of
-        ``aggregate_files``, comment at https://github.com/dask/dask/issues/9051".
+        WARNING: Passing a string argument to ``aggregate_files`` will result
+        in experimental behavior. This behavior may change in the future.
 
         Whether distinct file paths may be aggregated into the same output
         partition. This parameter is only used when `split_row_groups` is set to
@@ -416,14 +414,11 @@ def read_parquet(
             FutureWarning,
         )
 
-    # "Pre-deprecation" warning for `aggregate_files`
-    if aggregate_files:
+    # FutureWarning for `aggregate_files`
+    if aggregate_files and isinstance(aggregate_files, str):
         warnings.warn(
-            "The `aggregate_files` argument will be deprecated in the future. "
-            "Please consider using `from_map` to create a DataFrame collection "
-            "with a custom file-to-partition mapping.\n\n"
-            "If you strongly oppose the deprecation of `aggregate_files`, "
-            "please comment at https://github.com/dask/dask/issues/9051",
+            "String support for `aggregate_files` is experimental. "
+            "Behavior may change in the future. ",
             FutureWarning,
         )
 
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index 469835849..f45aeb241 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -19,7 +19,11 @@ from dask import config
 from dask.base import compute, compute_as_if_collection, is_dask_collection, tokenize
 from dask.dataframe import methods
 from dask.dataframe.core import DataFrame, Series, _Frame, map_partitions, new_dd_object
-from dask.dataframe.dispatch import group_split_dispatch, hash_object_dispatch
+from dask.dataframe.dispatch import (
+    group_split_dispatch,
+    hash_object_dispatch,
+    partd_encode_dispatch,
+)
 from dask.dataframe.utils import UNKNOWN_CATEGORIES
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import ShuffleLayer, SimpleShuffleLayer
@@ -521,16 +525,21 @@ class maybe_buffered_partd:
     If serialized, will return non-buffered partd. Otherwise returns a buffered partd
     """
 
-    def __init__(self, buffer=True, tempdir=None):
+    def __init__(self, encode_cls=None, buffer=True, tempdir=None):
         self.tempdir = tempdir or config.get("temporary_directory", None)
         self.buffer = buffer
         self.compression = config.get("dataframe.shuffle.compression", None)
+        self.encode_cls = encode_cls
+        if encode_cls is None:
+            import partd
+
+            self.encode_cls = partd.PandasBlocks
 
     def __reduce__(self):
         if self.tempdir:
-            return (maybe_buffered_partd, (False, self.tempdir))
+            return (maybe_buffered_partd, (self.encode_cls, False, self.tempdir))
         else:
-            return (maybe_buffered_partd, (False,))
+            return (maybe_buffered_partd, (self.encode_cls, False))
 
     def __call__(self, *args, **kwargs):
         import partd
@@ -556,9 +565,9 @@ class maybe_buffered_partd:
         if partd_compression:
             file = partd_compression(file)
         if self.buffer:
-            return partd.PandasBlocks(partd.Buffer(partd.Dict(), file))
+            return self.encode_cls(partd.Buffer(partd.Dict(), file))
         else:
-            return partd.PandasBlocks(file)
+            return self.encode_cls(file)
 
 
 def rearrange_by_column_disk(df, column, npartitions=None, compute=False):
@@ -577,7 +586,8 @@ def rearrange_by_column_disk(df, column, npartitions=None, compute=False):
     always_new_token = uuid.uuid1().hex
 
     p = ("zpartd-" + always_new_token,)
-    dsk1 = {p: (maybe_buffered_partd(),)}
+    encode_cls = partd_encode_dispatch(df._meta)
+    dsk1 = {p: (maybe_buffered_partd(encode_cls=encode_cls),)}
 
     # Partition data on disk
     name = "shuffle-partition-" + always_new_token
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
index 3c4cbec8d..d712fdea9 100644
--- a/dask/graph_manipulation.py
+++ b/dask/graph_manipulation.py
@@ -5,8 +5,8 @@ their inputs.
 from __future__ import annotations
 
 import uuid
-from collections.abc import Callable, Hashable, Set
-from typing import Any, Literal, TypeVar
+from collections.abc import Callable, Hashable
+from typing import Literal, TypeVar
 
 from dask.base import (
     clone_key,
@@ -321,7 +321,7 @@ def _bind_one(
 
     dsk = child.__dask_graph__()  # type: ignore
     new_layers: dict[str, Layer] = {}
-    new_deps: dict[str, Set[Any]] = {}
+    new_deps: dict[str, set[str]] = {}
 
     if isinstance(dsk, HighLevelGraph):
         try:
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index eb13ff013..0f13212ce 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -96,7 +96,7 @@ class Layer(Graph):
         return True
 
     @abc.abstractmethod
-    def get_output_keys(self) -> Set:
+    def get_output_keys(self) -> Set[Key]:
         """Return a set of all output keys
 
         Output keys are all keys in the layer that might be referenced by
@@ -405,16 +405,16 @@ class HighLevelGraph(Graph):
     """
 
     layers: Mapping[str, Layer]
-    dependencies: Mapping[str, Set[str]]
-    key_dependencies: dict[Key, Set[Key]]
+    dependencies: Mapping[str, set[str]]
+    key_dependencies: dict[Key, set[Key]]
     _to_dict: dict
     _all_external_keys: set
 
     def __init__(
         self,
         layers: Mapping[str, Graph],
-        dependencies: Mapping[str, Set[str]],
-        key_dependencies: dict[Key, Set[Key]] | None = None,
+        dependencies: Mapping[str, set[str]],
+        key_dependencies: dict[Key, set[Key]] | None = None,
     ):
         self.dependencies = dependencies
         self.key_dependencies = key_dependencies or {}
@@ -487,7 +487,7 @@ class HighLevelGraph(Graph):
             return cls._from_collection(name, layer, dependencies[0])
         layers = {name: layer}
         name_dep: set[str] = set()
-        deps: dict[str, Set[str]] = {name: name_dep}
+        deps: dict[str, set[str]] = {name: name_dep}
         for collection in toolz.unique(dependencies, key=id):
             if is_dask_collection(collection):
                 graph = collection.__dask_graph__()
@@ -583,7 +583,7 @@ class HighLevelGraph(Graph):
     def values(self) -> ValuesView[Any]:
         return self.to_dict().values()
 
-    def get_all_dependencies(self) -> dict[Key, Set[Key]]:
+    def get_all_dependencies(self) -> dict[Key, set[Key]]:
         """Get dependencies of all keys
 
         This will in most cases materialize all layers, which makes
@@ -616,7 +616,7 @@ class HighLevelGraph(Graph):
     @classmethod
     def merge(cls, *graphs: Graph) -> HighLevelGraph:
         layers: dict[str, Graph] = {}
-        dependencies: dict[str, Set[str]] = {}
+        dependencies: dict[str, set[str]] = {}
         for g in graphs:
             if isinstance(g, HighLevelGraph):
                 layers.update(g.layers)
diff --git a/dask/order.py b/dask/order.py
index 46e8c3145..9a8678660 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -78,12 +78,18 @@ Work towards *small goals* with *big steps*.
     good proxy for ordering.  This is usually a good idea and a sane default.
 """
 from collections import defaultdict, namedtuple
+from collections.abc import Mapping, MutableMapping
 from math import log
+from typing import Any, cast
 
-from dask.core import get_dependencies, get_deps, getcycle, reverse_dict
+from dask.core import get_dependencies, get_deps, getcycle, istask, reverse_dict
+from dask.typing import Key
 
 
-def order(dsk, dependencies=None):
+def order(
+    dsk: MutableMapping[Key, Any],
+    dependencies: MutableMapping[Key, set[Key]] | None = None,
+) -> dict[Key, int]:
     """Order nodes in dask graph
 
     This produces an ordering over our tasks that we use to break ties when
@@ -111,9 +117,10 @@ def order(dsk, dependencies=None):
     if not dsk:
         return {}
 
+    dsk = dict(dsk)
+
     if dependencies is None:
         dependencies = {k: get_dependencies(dsk, k) for k in dsk}
-
     dependents = reverse_dict(dependencies)
     num_needed, total_dependencies = ndependencies(dependencies, dependents)
     metrics = graph_metrics(dependencies, dependents, total_dependencies)
@@ -133,18 +140,29 @@ def order(dsk, dependencies=None):
     # tree, we skip processing it normally.
     # See https://github.com/dask/dask/issues/6745
     root_nodes = {k for k, v in dependents.items() if not v}
-    skip_root_node = len(root_nodes) == 1 and len(dsk) > 1
+    if len(root_nodes) > 1:
+        # This is also nice because it makes us robust to difference when
+        # computing vs persisting collections
+        root = cast(Key, object())
+
+        def _f(*args: Any, **kwargs: Any) -> None:
+            pass
 
+        dsk[root] = (_f, *root_nodes)
+        dependencies[root] = root_nodes
+        o = order(dsk, dependencies)
+        del o[root]
+        return o
+
+    init_stack: dict[Key, tuple] | set[Key] | list[Key]
     # Leaf nodes.  We choose one--the initial node--for each weakly connected subgraph.
     # Let's calculate the `initial_stack_key` as we determine `init_stack` set.
     init_stack = {
         # First prioritize large, tall groups, then prioritize the same as ``dependents_key``.
         key: (
             # at a high-level, work towards a large goal (and prefer tall and narrow)
-            -max_dependencies,
             num_dependents - max_heights,
             # tactically, finish small connected jobs first
-            min_dependencies,
             num_dependents - min_heights,  # prefer tall and narrow
             -total_dependents,  # take a big step
             # try to be memory efficient
@@ -154,8 +172,8 @@ def order(dsk, dependencies=None):
         )
         for key, num_dependents, (
             total_dependents,
-            min_dependencies,
-            max_dependencies,
+            _,
+            _,
             min_heights,
             max_heights,
         ) in (
@@ -168,41 +186,43 @@ def order(dsk, dependencies=None):
     # This value is static, so we pre-compute as the value of this dict.
     initial_stack_key = init_stack.__getitem__
 
-    def dependents_key(x):
+    def dependents_key(x: Key) -> tuple:
         """Choose a path from our starting task to our tactical goal
 
         This path is connected to a large goal, but focuses on completing
         a small goal and being memory efficient.
         """
+        assert dependencies is not None
+
         return (
             # Focus on being memory-efficient
             len(dependents[x]) - len(dependencies[x]) + num_needed[x],
+            # Do we favor deep or shallow branches?
+            #  -1: deep
+            #  +1: shallow
             -metrics[x][3],  # min_heights
             # tie-breaker
             StrComparable(x),
         )
 
-    def dependencies_key(x):
+    def dependencies_key(x: Key) -> tuple:
         """Choose which dependency to run as part of a reverse DFS
 
         This is very similar to both ``initial_stack_key``.
         """
+        assert dependencies is not None
         num_dependents = len(dependents[x])
         (
             total_dependents,
-            min_dependencies,
-            max_dependencies,
+            _,
+            _,
             min_heights,
             max_heights,
         ) = metrics[x]
         # Prefer short and narrow instead of tall in narrow, because we're going in
         # reverse along dependencies.
         return (
-            # at a high-level, work towards a large goal (and prefer short and narrow)
-            -max_dependencies,
             num_dependents + max_heights,
-            # tactically, finish small connected jobs first
-            min_dependencies,
             num_dependents + min_heights,  # prefer short and narrow
             -total_dependencies[x],  # go where the work is
             # try to be memory efficient
@@ -213,26 +233,23 @@ def order(dsk, dependencies=None):
             StrComparable(x),
         )
 
-    def finish_now_key(x):
-        """Determine the order of dependents that are ready to run and be released"""
-        return (-len(dependencies[x]), StrComparable(x))
-
+    root_total_dependencies = total_dependencies[list(root_nodes)[0]]
     # Computing this for all keys can sometimes be relatively expensive :(
     partition_keys = {
         key: (
-            (min_dependencies - total_dependencies[key] + 1)
+            (root_total_dependencies - total_dependencies[key] + 1)
             * (total_dependents - min_heights)
         )
         for key, (
             total_dependents,
-            min_dependencies,
+            _,
             _,
             min_heights,
             _,
         ) in metrics.items()
     }
 
-    result = {}
+    result: dict[Key, int] = {}
     i = 0
 
     # `inner_stack` is used to perform a DFS along dependencies.  Once emptied
@@ -247,7 +264,7 @@ def order(dsk, dependencies=None):
     # A "better path" is determined by comparing `partition_keys`.
     inner_stack = [min(init_stack, key=initial_stack_key)]
     inner_stack_pop = inner_stack.pop
-    inner_stacks = []
+    inner_stacks: list[list[Key]] = []
     inner_stacks_append = inner_stacks.append
     inner_stacks_extend = inner_stacks.extend
     inner_stacks_pop = inner_stacks.pop
@@ -258,13 +275,10 @@ def order(dsk, dependencies=None):
     # via `partition_key`.  A dependent goes to:
     #    1) `inner_stack` if it's better than our current target,
     #    2) `next_nodes` if the partition key is lower than it's parent,
-    #    3) `later_nodes` otherwise.
-    # When the inner stacks are depleted, we process `next_nodes`.  If `next_nodes` is
-    # empty (and `outer_stacks` is empty`), then we process `later_nodes` the same way.
+    # When the inner stacks are depleted, we process `next_nodes`.
     # These dicts use `partition_keys` as keys.  We process them by placing the values
     # in `outer_stack` so that the smallest keys will be processed first.
-    next_nodes = defaultdict(list)
-    later_nodes = defaultdict(list)
+    next_nodes: defaultdict[int, list[list[Key] | set[Key]]] = defaultdict(list)
 
     # `outer_stack` is used to populate `inner_stacks`.  From the time we partition the
     # dependents of a node, we group them: one list per partition key per parent node.
@@ -273,16 +287,13 @@ def order(dsk, dependencies=None):
     # partitioned, and we keep them in the order that we saw them (we will process them
     # in a FIFO manner).  By delaying sorting for as long as we can, we can first filter
     # out nodes that have already been computed.  All this complexity is worth it!
-    outer_stack = []
+    outer_stack: list[list[Key]] = []
     outer_stack_extend = outer_stack.extend
     outer_stack_pop = outer_stack.pop
 
     # Keep track of nodes that are in `inner_stack` or `inner_stacks` so we don't
     # process them again.
-    if skip_root_node:
-        seen = set(root_nodes)
-    else:
-        seen = set()  # seen in an inner_stack (and has dependencies)
+    seen = set(root_nodes)
     seen_update = seen.update
     seen_add = seen.add
 
@@ -324,10 +335,9 @@ def order(dsk, dependencies=None):
     # scheduler?  Should we defer to dynamic schedulers and let them behave like this
     # if they so choose?  Maybe.  However, I'm sensitive to the multithreaded scheduler,
     # which is heavily dependent on the ordering obtained here.
-    singles = {}
-    singles_items = singles.items()
+    singles: dict[Key, Key] = {}
     singles_clear = singles.clear
-    later_singles = []
+    later_singles: list[Key] = []
     later_singles_append = later_singles.append
     later_singles_clear = later_singles.clear
 
@@ -338,22 +348,24 @@ def order(dsk, dependencies=None):
     #   4. later_singles
     #   5. next_nodes
     #   6. outer_stack
-    #   7. later_nodes
-    #   8. init_stack
+    #   7. init_stack
 
     # alias for speed
     set_difference = set.difference
 
     is_init_sorted = False
+
     while True:
         while True:
             # Perform a DFS along dependencies until we complete our tactical goal
+            deps = set()
+            add_to_inner_stack = True
             if inner_stack:
                 item = inner_stack_pop()
                 if item in result:
                     continue
                 if num_needed[item]:
-                    if not skip_root_node or item not in root_nodes:
+                    if item not in root_nodes:
                         inner_stack.append(item)
                         deps = set_difference(dependencies[item], result)
                         if 1 < len(deps) < 1000:
@@ -365,6 +377,13 @@ def order(dsk, dependencies=None):
                         seen_update(deps)
                     if not singles:
                         continue
+                    # Only process singles once the inner_stack is fully
+                    # resolved. This is important because the singles path later
+                    # on verifies that running the single indeed opens an
+                    # opportunity to release soon by comparing the singles
+                    # parent's dependents with the inner_stack(s)
+                    if inner_stack and num_needed[inner_stack[-1]]:
+                        continue
                     process_singles = True
                 else:
                     result[item] = i
@@ -372,28 +391,6 @@ def order(dsk, dependencies=None):
                     deps = dependents[item]
                     add_to_inner_stack = True
 
-                    if metrics[item][3] == 1:  # min_height
-                        # Don't leave any dangling single nodes!  Finish all dependents that are
-                        # ready and are also root nodes.
-                        finish_now = {
-                            dep
-                            for dep in deps
-                            if not dependents[dep] and num_needed[dep] == 1
-                        }
-                        if finish_now:
-                            deps -= finish_now  # Safe to mutate
-                            if len(finish_now) > 1:
-                                finish_now = sorted(finish_now, key=finish_now_key)
-                            for dep in finish_now:
-                                result[dep] = i
-                                i += 1
-                            add_to_inner_stack = False
-                        elif skip_root_node:
-                            for dep in root_nodes:
-                                num_needed[dep] -= 1
-                                # Use remove here to complain loudly if our assumptions change
-                                deps.remove(dep)  # Safe to mutate
-
                     if deps:
                         for dep in deps:
                             num_needed[dep] -= 1
@@ -409,45 +406,25 @@ def order(dsk, dependencies=None):
             elif later_singles:
                 # No need to be optimistic: all nodes in `later_singles` will free a dependency
                 # when run, so no need to check whether dependents are in `seen`.
-                deps = set()
                 for single in later_singles:
                     if single in result:
                         continue
                     while True:
-                        dep2 = dependents[single]
+                        deps_singles = dependents[single]
                         result[single] = i
                         i += 1
-                        if metrics[single][3] == 1:  # min_height
-                            # Don't leave any dangling single nodes!  Finish all dependents that are
-                            # ready and are also root nodes.
-                            finish_now = {
-                                dep
-                                for dep in dep2
-                                if not dependents[dep] and num_needed[dep] == 1
-                            }
-                            if finish_now:
-                                dep2 -= finish_now  # Safe to mutate
-                                if len(finish_now) > 1:
-                                    finish_now = sorted(finish_now, key=finish_now_key)
-                                for dep in finish_now:
-                                    result[dep] = i
-                                    i += 1
-                            elif skip_root_node:
-                                for dep in root_nodes:
-                                    num_needed[dep] -= 1
-                                    # Use remove here to complain loudly if our assumptions change
-                                    dep2.remove(dep)  # Safe to mutate
-                        if dep2:
-                            for dep in dep2:
+                        if deps_singles:
+                            for dep in deps_singles:
                                 num_needed[dep] -= 1
-                            if len(dep2) == 1:
+                            if len(deps_singles) == 1:
                                 # Fast path!  We trim down `dep2` above hoping to reach here.
-                                (single,) = dep2
+                                (single,) = deps_singles
                                 if not num_needed[single]:
                                     # Keep it going!
-                                    dep2 = dependents[single]
+                                    deps_singles = dependents[single]
                                     continue
-                            deps |= dep2
+                            deps |= deps_singles
+                        del deps_singles
                         break
                 later_singles_clear()
                 deps = set_difference(deps, result)
@@ -460,74 +437,71 @@ def order(dsk, dependencies=None):
 
             if process_singles and singles:
                 # We gather all dependents of all singles into `deps`, which we then process below.
-                # A lingering question is: what should we use for `item`?  `item_key` is used to
-                # determine whether each dependent goes to `next_nodes` or `later_nodes`.  Currently,
-                # we use the last value of `item` (i.e., we don't do anything).
-                deps = set()
+
                 add_to_inner_stack = True if inner_stack or inner_stacks else False
-                for single, parent in singles_items:
-                    if single in result:
-                        continue
+                singles_keys = set_difference(set(singles), result)
+
+                # NOTE: If this was too slow, LIFO would be a decent
+                # approximation
+                for single in sorted(singles_keys, key=lambda x: partition_keys[x]):
+                    # We want to run the singles if they are either releasing a
+                    # dependency directly or that they may be releasing a
+                    # dependency once the current critical path / inner_stack is
+                    # walked.
+                    # By using `seen` here this is more permissive since it also
+                    # includes tasks in a future critical path / inner_stacks
+                    # but it would require additional state to make this
+                    # distinction and we don't have enough data to dermine if
+                    # this is worth it.
+                    parent = singles[single]
                     if (
-                        add_to_inner_stack
-                        and len(set_difference(dependents[parent], result)) > 1
+                        len(
+                            set_difference(
+                                set_difference(dependents[parent], result),
+                                seen,
+                            )
+                        )
+                        > 1
                     ):
                         later_singles_append(single)
                         continue
-
                     while True:
-                        dep2 = dependents[single]
+                        deps_singles = dependents[single]
                         result[single] = i
                         i += 1
-                        if metrics[single][3] == 1:  # min_height
-                            # Don't leave any dangling single nodes!  Finish all dependents that are
-                            # ready and are also root nodes.
-                            finish_now = {
-                                dep
-                                for dep in dep2
-                                if not dependents[dep] and num_needed[dep] == 1
-                            }
-                            if finish_now:
-                                dep2 -= finish_now  # Safe to mutate
-                                if len(finish_now) > 1:
-                                    finish_now = sorted(finish_now, key=finish_now_key)
-                                for dep in finish_now:
-                                    result[dep] = i
-                                    i += 1
-                            elif skip_root_node:
-                                for dep in root_nodes:
-                                    num_needed[dep] -= 1
-                                    # Use remove here to complain loudly if our assumptions change
-                                    dep2.remove(dep)  # Safe to mutate
-                        if dep2:
-                            for dep in dep2:
+                        if deps_singles:
+                            for dep in deps_singles:
                                 num_needed[dep] -= 1
                             if add_to_inner_stack:
-                                already_seen = dep2 & seen
+                                already_seen = deps_singles & seen
                                 if already_seen:
-                                    if len(dep2) == len(already_seen):
+                                    # This means that the singles path also
+                                    # leads to the current or previous strategic
+                                    # path
+                                    if len(deps_singles) == len(already_seen):
                                         if len(already_seen) == 1:
                                             (single,) = already_seen
                                             if not num_needed[single]:
-                                                dep2 = dependents[single]
+                                                deps_singles = dependents[single]
                                                 continue
                                         break
-                                    dep2 = dep2 - already_seen
+                                    deps_singles = deps_singles - already_seen
                             else:
-                                already_seen = False
-                            if len(dep2) == 1:
+                                already_seen = set()
+                            if len(deps_singles) == 1:
                                 # Fast path!  We trim down `dep2` above hoping to reach here.
-                                (single,) = dep2
+                                (single,) = deps_singles
                                 if not num_needed[single]:
                                     if not already_seen:
                                         # Keep it going!
-                                        dep2 = dependents[single]
+                                        deps_singles = dependents[single]
                                         continue
                                     later_singles_append(single)
                                     break
-                            deps |= dep2
+                            deps |= deps_singles
+                        del deps_singles
                         break
-
+                del singles_keys
                 deps = set_difference(deps, result)
                 singles_clear()
                 if not deps:
@@ -535,8 +509,7 @@ def order(dsk, dependencies=None):
                 add_to_inner_stack = False
 
             # If inner_stack is empty, then we typically add the best dependent to it.
-            # However, we don't add to it if we complete a node early via "finish_now" above
-            # or if a dependent is already on an inner_stack.  In this case, we add the
+            # However, we don't add to it if a dependent is already on an inner_stack.  In this case, we add the
             # dependents (not in an inner_stack) to next_nodes or later_nodes to handle later.
             # This serves three purposes:
             #   1. shrink `deps` so that it can be processed faster,
@@ -549,6 +522,7 @@ def order(dsk, dependencies=None):
                         (dep,) = already_seen
                         if not num_needed[dep]:
                             singles[dep] = item
+                        del dep
                     continue
                 add_to_inner_stack = False
                 deps = deps - already_seen
@@ -556,30 +530,19 @@ def order(dsk, dependencies=None):
             if len(deps) == 1:
                 # Fast path!  We trim down `deps` above hoping to reach here.
                 (dep,) = deps
-                if not inner_stack:
-                    if add_to_inner_stack:
-                        inner_stack = [dep]
-                        inner_stack_pop = inner_stack.pop
-                        seen_add(dep)
-                        continue
-                    key = partition_keys[dep]
-                else:
-                    key = partition_keys[dep]
-                    if key < partition_keys[inner_stack[0]]:
-                        # Run before `inner_stack` (change tactical goal!)
-                        inner_stacks_append(inner_stack)
-                        inner_stack = [dep]
-                        inner_stack_pop = inner_stack.pop
-                        seen_add(dep)
-                        continue
+                if add_to_inner_stack and not inner_stack:
+                    inner_stack = [dep]
+                    inner_stack_pop = inner_stack.pop
+                    seen_add(dep)
+                    continue
+                key = partition_keys[dep]
                 if not num_needed[dep]:
                     # We didn't put the single dependency on the stack, but we should still
                     # run it soon, because doing so may free its parent.
                     singles[dep] = item
-                elif key < partition_keys[item]:
-                    next_nodes[key].append(deps)
                 else:
-                    later_nodes[key].append(deps)
+                    next_nodes[key].append(deps)
+                del dep, key
             elif len(deps) == 2:
                 # We special-case when len(deps) == 2 so that we may place a dep on singles.
                 # Otherwise, the logic here is the same as when `len(deps) > 2` below.
@@ -620,93 +583,88 @@ def order(dsk, dependencies=None):
                                 later_singles_append(dep2)
                             else:
                                 singles[dep2] = item
-                        elif key2 < partition_keys[item]:
-                            next_nodes[key2].append([dep2])
                         else:
-                            later_nodes[key2].append([dep2])
+                            next_nodes[key2].append([dep2])
                     else:
                         item_key = partition_keys[item]
-                        if key2 < item_key:
-                            next_nodes[key].append([dep])
-                            next_nodes[key2].append([dep2])
-                        elif key < item_key:
-                            next_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
-                        else:
-                            later_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
+                        for k, d in [(key, dep), (key2, dep2)]:
+                            if not num_needed[d]:
+                                if process_singles:
+                                    later_singles_append(d)
+                                else:
+                                    singles[d] = item
+                            else:
+                                next_nodes[k].append([d])
+                        del item_key
+                    del prev_key
                 else:
+                    assert not inner_stack
                     if add_to_inner_stack:
+                        inner_stack = [dep]
+                        inner_stack_pop = inner_stack.pop
+                        seen_add(dep)
                         if not num_needed[dep2]:
-                            inner_stacks_append(inner_stack)
-                            inner_stack = [dep]
-                            inner_stack_pop = inner_stack.pop
-                            seen_add(dep)
                             singles[dep2] = item
                         elif key == key2 and 5 * partition_keys[item] > 22 * key:
-                            inner_stacks_append(inner_stack)
                             inner_stacks_append([dep2])
-                            inner_stack = [dep]
-                            inner_stack_pop = inner_stack.pop
-                            seen_update(deps)
+                            seen_add(dep2)
                         else:
-                            inner_stacks_append(inner_stack)
-                            inner_stack = [dep]
-                            inner_stack_pop = inner_stack.pop
-                            seen_add(dep)
-                            if key2 < partition_keys[item]:
-                                next_nodes[key2].append([dep2])
-                            else:
-                                later_nodes[key2].append([dep2])
-                    else:
-                        item_key = partition_keys[item]
-                        if key2 < item_key:
-                            next_nodes[key].append([dep])
                             next_nodes[key2].append([dep2])
-                        elif key < item_key:
-                            next_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
-                        else:
-                            later_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
+                    else:
+                        for k, d in [(key, dep), (key2, dep2)]:
+                            next_nodes[k].append([d])
+                del dep, dep2, key, key2
             else:
                 # Slow path :(.  This requires grouping by partition_key.
-                dep_pools = defaultdict(list)
+                dep_pools = defaultdict(set)
+                possible_singles = defaultdict(set)
                 for dep in deps:
-                    dep_pools[partition_keys[dep]].append(dep)
+                    pkey = partition_keys[dep]
+                    if not num_needed[dep] and not process_singles:
+                        possible_singles[pkey].add(dep)
+                    dep_pools[pkey].add(dep)
                 item_key = partition_keys[item]
                 if inner_stack:
                     # If we have an inner_stack, we need to look for a "better" path
                     prev_key = partition_keys[inner_stack[0]]
                     now_keys = []  # < inner_stack[0]
+                    psingles = set()
                     for key, vals in dep_pools.items():
                         if key < prev_key:
                             now_keys.append(key)
-                        elif key < item_key:
-                            next_nodes[key].append(vals)
                         else:
-                            later_nodes[key].append(vals)
+                            psingles = possible_singles[key]
+                            for s in psingles:
+                                singles[s] = item
+                            vals -= psingles
+                            next_nodes[key].append(vals)
+                        del vals, key
+                    del psingles
                     if now_keys:
                         # Run before `inner_stack` (change tactical goal!)
                         inner_stacks_append(inner_stack)
                         if 1 < len(now_keys):
                             now_keys.sort(reverse=True)
                         for key in now_keys:
+                            pool: set[Key] | list[Key]
                             pool = dep_pools[key]
                             if 1 < len(pool) < 100:
-                                pool.sort(key=dependents_key, reverse=True)
+                                pool = sorted(pool, key=dependents_key, reverse=True)
                             inner_stacks_extend([dep] for dep in pool)
                             seen_update(pool)
+                            del pool
                         inner_stack = inner_stacks_pop()
                         inner_stack_pop = inner_stack.pop
+                    del now_keys, prev_key
                 else:
                     # If we don't have an inner_stack, then we don't need to look
                     # for a "better" path, but we do need traverse along dependents.
                     if add_to_inner_stack:
+                        min_pool: list[Key] | set[Key]
                         min_key = min(dep_pools)
                         min_pool = dep_pools.pop(min_key)
                         if len(min_pool) == 1:
-                            inner_stack = min_pool
+                            inner_stack = list(min_pool)
                             seen_update(inner_stack)
                         elif (
                             10 * item_key > 11 * len(min_pool) * len(min_pool) * min_key
@@ -721,7 +679,9 @@ def order(dsk, dependencies=None):
                             # what we have easily available.  It is obviously very specific to our
                             # choice of partition_key.  Dask tests take this route about 40%.
                             if len(min_pool) < 100:
-                                min_pool.sort(key=dependents_key, reverse=True)
+                                min_pool = sorted(
+                                    min_pool, key=dependents_key, reverse=True
+                                )
                             inner_stacks_extend([dep] for dep in min_pool)
                             inner_stack = inner_stacks_pop()
                             seen_update(min_pool)
@@ -733,13 +693,15 @@ def order(dsk, dependencies=None):
                                 inner_stack = [min_pool.pop()]
                             next_nodes[min_key].append(min_pool)
                             seen_update(inner_stack)
-
+                        del min_pool, min_key
                         inner_stack_pop = inner_stack.pop
                     for key, vals in dep_pools.items():
-                        if key < item_key:
-                            next_nodes[key].append(vals)
-                        else:
-                            later_nodes[key].append(vals)
+                        psingles = possible_singles[key]
+                        for s in psingles:
+                            singles[s] = item
+                        vals -= psingles
+                        next_nodes[key].append(vals)
+                        del key, vals
 
         if len(dependencies) == len(result):
             break  # all done!
@@ -749,28 +711,24 @@ def order(dsk, dependencies=None):
                 # `outer_stacks` may not be empty here--it has data from previous `next_nodes`.
                 # Since we pop things off of it (onto `inner_nodes`), this means we handle
                 # multiple `next_nodes` in a LIFO manner.
-                outer_stack_extend(reversed(next_nodes[key]))
-            next_nodes = defaultdict(list)
+                outer_stack_extend(list(el) for el in reversed(next_nodes[key]))
+            next_nodes.clear()
 
+        outer_deps = []
         while outer_stack:
             # Try to add a few items to `inner_stacks`
-            deps = [x for x in outer_stack_pop() if x not in result]
-            if deps:
-                if 1 < len(deps) < 100:
-                    deps.sort(key=dependents_key, reverse=True)
-                inner_stacks_extend([dep] for dep in deps)
-                seen_update(deps)
+            outer_deps = [x for x in outer_stack_pop() if x not in result]
+            if outer_deps:
+                if 1 < len(outer_deps) < 100:
+                    outer_deps.sort(key=dependents_key, reverse=True)
+                inner_stacks_extend([dep] for dep in outer_deps)
+                seen_update(outer_deps)
                 break
+        del outer_deps
 
         if inner_stacks:
             continue
 
-        if later_nodes:
-            # You know all those dependents with large keys we've been hanging onto to run "later"?
-            # Well, "later" has finally come.
-            next_nodes, later_nodes = later_nodes, next_nodes
-            continue
-
         # We just finished computing a connected group.
         # Let's choose the first `item` in the next group to compute.
         # If we have few large groups left, then it's best to find `item` by taking a minimum.
@@ -778,8 +736,7 @@ def order(dsk, dependencies=None):
         # If we have many tiny groups left, then it's best to simply iterate.
         if not is_init_sorted:
             prev_len = len(init_stack)
-            if type(init_stack) is dict:
-                init_stack = set(init_stack)
+            init_stack = set(init_stack)
             init_stack = set_difference(init_stack, result)
             N = len(init_stack)
             m = prev_len - N
@@ -795,7 +752,7 @@ def order(dsk, dependencies=None):
             init_stack_pop = init_stack.pop
             is_init_sorted = True
 
-        if skip_root_node and item in root_nodes:
+        if item in root_nodes:
             item = init_stack_pop()
 
         while item in result:
@@ -805,7 +762,11 @@ def order(dsk, dependencies=None):
     return result
 
 
-def graph_metrics(dependencies, dependents, total_dependencies):
+def graph_metrics(
+    dependencies: Mapping[Key, set[Key]],
+    dependents: Mapping[Key, set[Key]],
+    total_dependencies: Mapping[Key, int],
+) -> dict[Key, tuple[int, int, int, int, int]]:
     r"""Useful measures of a graph used by ``dask.order.order``
 
     Example DAG (a1 has no dependencies; b2 and c1 are root nodes):
@@ -819,8 +780,10 @@ def graph_metrics(dependencies, dependents, total_dependencies):
     For each key we return:
 
     1.  **total_dependents**: The number of keys that can only be run
-        after this key is run.  The root nodes have value 1 while deep child
-        nodes will have larger values.
+        after this key is run.
+        Note that this is only exact for trees. (undirected) cycles will cause
+        double counting of nodes. Therefore, this metric is an upper bound
+        approximation.
 
         1
         |
@@ -882,7 +845,7 @@ def graph_metrics(dependencies, dependents, total_dependencies):
     """
     result = {}
     num_needed = {k: len(v) for k, v in dependents.items() if v}
-    current = []
+    current: list[Key] = []
     current_pop = current.pop
     current_append = current.append
     for key, deps in dependents.items():
@@ -915,18 +878,18 @@ def graph_metrics(dependencies, dependents, total_dependencies):
             )
         else:
             (
-                total_dependents,
-                min_dependencies,
-                max_dependencies,
-                min_heights,
-                max_heights,
+                total_dependents_,
+                min_dependencies_,
+                max_dependencies_,
+                min_heights_,
+                max_heights_,
             ) = zip(*(result[parent] for parent in dependents[key]))
             result[key] = (
-                1 + sum(total_dependents),
-                min(min_dependencies),
-                max(max_dependencies),
-                1 + min(min_heights),
-                1 + max(max_heights),
+                1 + sum(total_dependents_),
+                min(min_dependencies_),
+                max(max_dependencies_),
+                1 + min(min_heights_),
+                1 + max(max_heights_),
             )
         for child in dependencies[key]:
             num_needed[child] -= 1
@@ -935,7 +898,9 @@ def graph_metrics(dependencies, dependents, total_dependencies):
     return result
 
 
-def ndependencies(dependencies, dependents):
+def ndependencies(
+    dependencies: Mapping[Key, set[Key]], dependents: Mapping[Key, set[Key]]
+) -> tuple[dict[Key, int], dict[Key, int]]:
     """Number of total data elements on which this key depends
 
     For each key we return the number of tasks that must be run for us to run
@@ -963,7 +928,7 @@ def ndependencies(dependencies, dependents):
             result[k] = 1
 
     num_dependencies = num_needed.copy()
-    current = []
+    current: list[Key] = []
     current_pop = current.pop
     current_append = current.append
 
@@ -1001,10 +966,12 @@ class StrComparable:
 
     __slots__ = ("obj",)
 
-    def __init__(self, obj):
+    obj: Any
+
+    def __init__(self, obj: Any):
         self.obj = obj
 
-    def __lt__(self, other):
+    def __lt__(self, other: Any) -> bool:
         try:
             return self.obj < other.obj
         except Exception:
@@ -1023,7 +990,11 @@ OrderInfo = namedtuple(
 )
 
 
-def diagnostics(dsk, o=None, dependencies=None):
+def diagnostics(
+    dsk: MutableMapping[Key, Any],
+    o: Mapping[Key, int] | None = None,
+    dependencies: MutableMapping[Key, set[Key]] | None = None,
+) -> tuple[dict[Key, OrderInfo], list[int]]:
     """Simulate runtime metrics as though running tasks one at a time in order.
 
     These diagnostics can help reveal behaviors of and issues with ``order``.
@@ -1076,3 +1047,40 @@ def diagnostics(dsk, o=None, dependencies=None):
         for key, val in o.items()
     }
     return rv, pressure
+
+
+def _f() -> None:
+    ...
+
+
+def _convert_task(task: Any) -> Any:
+    if istask(task):
+        assert callable(task[0])
+        new_spec: list[Any] = []
+        for el in task[1:]:
+            if isinstance(el, (str, int)):
+                new_spec.append(el)
+            elif isinstance(el, tuple):
+                if istask(el):
+                    new_spec.append(_convert_task(el))
+                else:
+                    new_spec.append(el)
+            elif isinstance(el, list):
+                new_spec.append([_convert_task(e) for e in el])
+        return (_f, *new_spec)
+    else:
+        return task
+
+
+def sanitize_dsk(dsk: MutableMapping[Key, Any]) -> dict:
+    """Take a dask graph and replace callables with a dummy function and remove
+    payload data like numpy arrays, dataframes, etc.
+    """
+    new = {}
+    for key, values in dsk.items():
+        new_key = key
+        new[new_key] = _convert_task(values)
+    if get_deps(new) != get_deps(dsk):
+        # The switch statement in _convert likely dropped some keys
+        raise RuntimeError("Sanitization failed to preserve topology.")
+    return new
diff --git a/dask/typing.py b/dask/typing.py
index 5cb71bd75..5a47f2d85 100644
--- a/dask/typing.py
+++ b/dask/typing.py
@@ -26,6 +26,8 @@ PostComputeCallable = Callable
 
 
 Key: TypeAlias = Union[str, bytes, int, float, tuple["Key", ...]]
+# FIXME: This type is a little misleading. Low level graphs are often
+# MutableMappings but HLGs are not
 Graph: TypeAlias = Mapping[Key, Any]
 # Potentially nested list of Dask keys
 NestedKeys: TypeAlias = list[Union[Key, "NestedKeys"]]
diff --git a/docs/source/dataframe-hive.rst b/docs/source/dataframe-hive.rst
new file mode 100644
index 000000000..8b8c2ba80
--- /dev/null
+++ b/docs/source/dataframe-hive.rst
@@ -0,0 +1,247 @@
+.. _dataframe.hive:
+
+Using Hive Partitioning with Dask
+=================================
+
+.. currentmodule:: dask.dataframe
+
+It is sometimes useful to write your dataset with a hive-like directory scheme.
+For example, if your dataframe contains ``'year'`` and ``'semester'`` columns,
+a hive-based directory structure might look something like the following::
+
+    output-path/
+    ├── year=2022/
+    │   ├── semester=fall/
+    │   │   └── part.0.parquet
+    │   └── semester=spring/
+    │       ├── part.0.parquet
+    │       └── part.1.parquet
+    └── year=2023/
+        └── semester=fall/
+            └── part.1.parquet
+
+The use of this self-describing structure implies that all rows within
+the ``'output-path/year=2022/semester=fall/'`` directory will contain
+the value ``2022`` in the ``'year'`` column and the value ``'fall'``
+in the ``'semester'`` column.
+
+The primary advantage of generating a hive-partitioned dataset
+is that certain IO filters can be applied by :func:`read_parquet`
+without the need to parse any file metadata. In other words,
+the following command will typically be faster when the dataset
+is already hive-partitioned on the ``'year'`` column.
+
+.. code-block:: python
+
+    >>> dd.read_parquet("output-path", filters=[("year", ">", 2022)])
+
+Writing Parquet Data with Hive Partitioning
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Dask's :func:`to_parquet` function will produce a hive-partitioned
+directory scheme automatically when the ``partition_on`` option is used.
+
+.. code-block:: python
+
+    >>> df.to_parquet("output-path", partition_on=["year", "semester"])
+
+    >>> os.listdir("output-path")
+    ["year=2022", "year=2023"]
+
+    >>> os.listdir("output-path/year=2022")
+    ["semester=fall", "semester=spring"]
+
+    >>> os.listdir("output-path/year=2022/semester=spring")
+    ['part.0.parquet', 'part.1.parquet']
+
+
+It is important to recognize that Dask will **not** aggregate the
+data files written within each of the leaf directories. This is
+because each of the DataFrame partitions is written independently 
+during the execution of the :func:`to_parquet` task graph. In order 
+to write out data for partition `i`, the partition-i write task will
+perform a `groupby` operation on columns ``["year", "semester"]``,
+and then each distinct group will be written to the corresponding
+directory using the file name ``'part.{i}.parquet'``. Therefore, it
+is possible for a hive-partitioned write to produce a large number
+of files in every leaf directory (one for each DataFrame partition).
+
+If your application requires you to produce a single parquet file
+for each hive partition, one possible solution is to sort or shuffle
+on the partitioning columns before calling :func:`to_parquet`.
+
+.. code-block:: python
+
+    >>> partition_on = ["year", "semester"]
+
+    >>> df.shuffle(on=partition_on).to_parquet(partition_on=partition_on)
+
+Using a global shuffle like this is extremely expensive, and should be
+avoided whenever possible. However, it is also guaranteed to produce
+the minimum number of files, which may be worth the sacrifice at times.
+
+
+Reading Parquet Data with Hive Partitioning
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+In most cases, :func:`read_parquet` will process hive-partitioned
+data automatically. By default, all hive-partitioned columns will
+be interpreted as categorical columns.
+
+.. code-block:: python
+
+    >>> ddf = dd.read_parquet("output-path", columns=["year", "semester"])
+
+    >>> ddf
+    Dask DataFrame Structure:
+                            year         semester
+    npartitions=4                                  
+                category[known]  category[known]
+                            ...              ...
+                            ...              ...
+                            ...              ...
+                            ...              ...
+    Dask Name: read-parquet, 1 graph layer
+
+    >>> ddf.compute()
+    year semester
+    0  2022     fall
+    1  2022     fall
+    2  2022     fall
+    3  2022   spring
+    4  2022   spring
+    5  2022   spring
+    6  2023     fall
+    7  2023     fall
+
+
+Defining a Custom Partitioning Schema
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When utilizing ``engine='pyarrow'``, it is possible to specify a custom
+schema for the hive-partitioned columns. The columns will then be read
+using the specified types and not as `category`.
+
+.. code-block:: python
+
+    >>> schema = pa.schema([("year", pa.int16()), ("semester", pa.string())])
+
+    >>> ddf2 = dd.read_parquet(
+    ...     path,
+    ...     columns=["year", "semester"],
+    ...     dataset={"partitioning": {"flavor": "hive", "schema": schema}}
+    ... )
+    Dask DataFrame Structure:
+                    year semester
+    npartitions=4                
+                int16   object
+                    ...      ...
+                    ...      ...
+                    ...      ...
+                    ...      ...
+
+
+If any of your hive-partitioned columns contain null values, you
+**must** specify the partitioning schema in this way.
+
+Although it is not required, we also recommend that you specify
+the partitioning schema if you need to partition on high-cardinality
+columns. This is because the default ``'category'`` dtype will
+track the known categories in a way that can significantly increase
+the overall memory footprint of your Dask collection. In fact,
+:func:`read_parquet` already clears the "known categories" of other
+columns for this same reason (see :doc:`dataframe-categoricals`).
+
+
+Best Practices
+--------------
+
+Although hive partitioning can sometimes improve read performance
+by simplifying filtering, it can also lead to degraded performance
+and errors in other cases.
+
+
+Avoid High Cardinality
+~~~~~~~~~~~~~~~~~~~~~~
+
+A good rule of thumb is to avoid partitioning on `float` columns,
+or any column containing many unique values (i.e. high cardinality).
+
+The most common cause of poor user experience with hive partitioning
+is high-cardinality of the partitioning column(s). For example, if 
+you try to partition on a column with millions of unique values, then
+`:func:`to_parquet`` will need to generate millions of directories. The
+management of these directories is likely to put strain on the file
+system, and the need for many small files within each directory
+is sure to compound the issue.
+
+
+Use Simple Data Types for Partitioning
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Since hive-partitioned data is "self describing," we suggest that
+you avoid partitioning on complex data types, and opt for integer
+or string-based data types whenever possible. If your data type
+cannot be easily inferred from the string value used to define the
+directory name, then the IO engine may struggle to parse the values.
+
+For example, directly partitioning on a column with a ``datetime64``
+dtype might produce a directory name like the following::
+
+    output-path/
+    ├── date=2022-01-01 00:00:00/
+    ├── date=2022-02-01 00:00:00/
+    ├── ...
+    └── date=2022-12-01 00:00:00/
+
+These directory names will not be correctly interpreted as
+``datetime64`` values, and are even considered illegal on Windows
+systems. For more-reliable behavior, we recommend that such a column
+be decomposed into one or more "simple" columns. For example, one
+could easily use ``'date'`` to construct ``'year'``, ``'month'``,
+and ``'day'`` columns (as needed).
+
+
+Aggregate Files at Read Time
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. warning::
+    The ``aggregate_files`` argument is currently listed as
+    experimental. However, there are currently no plans to remove
+    the argument or change it's behavior in a future release.
+
+Since hive-partitioning will typically produce a large number of
+small files, :func:`read_parquet` performance will usually benefit
+from proper usage of the ``aggregate_files`` argument. Take the
+following dataset for example::
+
+    dataset-path/
+    ├── region=1/
+    │   ├── section=a/
+    │   │   └── 01.parquet
+    │   │   └── 02.parquet
+    │   │   └── 03.parquet
+    │   ├── section=b/
+    │   └── └── 04.parquet
+    │   └── └── 05.parquet
+    └── region=2/
+        ├── section=a/
+        │   ├── 06.parquet
+        │   ├── 07.parquet
+        │   ├── 08.parquet
+
+If we set ``aggregate_files=True`` for this case, we are telling Dask
+that any of the parquet data files may be aggregated into the same output
+DataFrame partition. If, instead, we specify the name of a partitioning
+column (e.g. ``'region'`` or ``'section'``), we allow the aggregation of
+any two files sharing a file path up to, and including, the corresponding
+directory name. For example, if ``aggregate_files`` is set to ``'section'``,
+``04.parquet`` and ``05.parquet`` may be aggregated together, but
+``03.parquet`` and ``04.parquet`` cannot be. If, however, ``aggregate_files``
+is set to ``'region'``, ``04.parquet`` may be aggregated with ``05.parquet``,
+**and** ``03.parquet`` may be aggregated with ``04.parquet``.
+
+Using ``aggregate_files`` will typically improve performance by making it
+more likely for DataFrame partitions to approach the size specified by the
+``blocksize`` argument. In contrast, default behavior may produce a large
+number of partitions that are much smaller than ``blocksize``.
diff --git a/docs/source/dataframe-parquet.rst b/docs/source/dataframe-parquet.rst
index e514269b5..813050ba9 100644
--- a/docs/source/dataframe-parquet.rst
+++ b/docs/source/dataframe-parquet.rst
@@ -260,6 +260,7 @@ moderate dataset sizes.
 File Names
 ~~~~~~~~~~
 
+Unless the `partition_on` option is used (see :doc:`dataframe-hive`),
 :func:`to_parquet` will write one file per Dask dataframe partition to the
 output directory. By default these files will have names like
 ``part.0.parquet``, ``part.1.parquet``, etc. If you wish to alter this naming
@@ -279,6 +280,15 @@ their partition indices.
     >>> os.listdir("/path/to/parquet")
     ["data-0.parquet", "data-1.parquet", "data-2.parquet"]
 
+Hive Partitioning
+~~~~~~~~~~~~~~~~~
+
+It is sometimes useful to write a parquet dataset with a hive-like directory scheme
+(e.g. ``'/year=2022/month=12/day=25'``). :func:`to_parquet` will automatically
+produce a dataset with this kind of directory structure when the ``partition_on``
+option is used. In most cases, :func:`to_parquet` will handle hive partitioning
+automatically. See :doc:`dataframe-hive` for more information.
+
 .. _parquet: https://parquet.apache.org/
 .. _glob string: https://docs.python.org/3/library/glob.html
 .. _fsspec: https://filesystem-spec.readthedocs.io
diff --git a/docs/source/dataframe.rst b/docs/source/dataframe.rst
index 6c3f45996..1278411c6 100644
--- a/docs/source/dataframe.rst
+++ b/docs/source/dataframe.rst
@@ -17,6 +17,7 @@ Dask DataFrame
    dataframe-categoricals.rst
    dataframe-extend.rst
    dataframe-parquet.rst
+   dataframe-hive.rst
    dataframe-sql.rst
    dataframe-api.rst
 
diff --git a/pyproject.toml b/pyproject.toml
index c29035253..e938d7aa3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -166,6 +166,14 @@ warn_redundant_casts = true
 warn_unused_ignores = true
 warn_unreachable = true
 
+[[tool.mypy.overrides]]
+
+# Recent or recently overhauled modules featuring stricter validation
+module = [
+    "dask.order",
+]
+allow_untyped_defs = false
+
 [tool.codespell]
 ignore-words-list = "coo,nd"
 skip = "docs/source/changelog.rst"


Test patch:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index 6c53cd6cb..a4ea31d11 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -2861,13 +2861,12 @@ def test_split_row_groups_int_aggregate_files(tmpdir, engine, split_row_groups):
 
     # Read back with both `split_row_groups>1` and
     # `aggregate_files=True`
-    with pytest.warns(FutureWarning, match="argument will be deprecated"):
-        ddf2 = dd.read_parquet(
-            str(tmpdir),
-            engine=engine,
-            split_row_groups=split_row_groups,
-            aggregate_files=True,
-        )
+    ddf2 = dd.read_parquet(
+        str(tmpdir),
+        engine=engine,
+        split_row_groups=split_row_groups,
+        aggregate_files=True,
+    )
 
     # Check that we are aggregating files as expected
     npartitions_expected = math.ceil((size / row_group_size) / split_row_groups)
@@ -3052,13 +3051,23 @@ def test_split_adaptive_files(
         write_index=False,
     )
 
-    with pytest.warns(FutureWarning, match="argument will be deprecated"):
+    aggregate_files = partition_on if partition_on else True
+    if isinstance(aggregate_files, str):
+        with pytest.warns(FutureWarning, match="Behavior may change"):
+            ddf2 = dd.read_parquet(
+                str(tmpdir),
+                engine=read_engine,
+                blocksize=blocksize,
+                split_row_groups="adaptive",
+                aggregate_files=aggregate_files,
+            )
+    else:
         ddf2 = dd.read_parquet(
             str(tmpdir),
             engine=read_engine,
             blocksize=blocksize,
             split_row_groups="adaptive",
-            aggregate_files=partition_on if partition_on else True,
+            aggregate_files=aggregate_files,
         )
 
     # Check that files where aggregated as expected
@@ -3103,7 +3112,7 @@ def test_split_adaptive_aggregate_files(
         partition_on=partition_on,
         write_index=False,
     )
-    with pytest.warns(FutureWarning, match="argument will be deprecated"):
+    with pytest.warns(FutureWarning, match="Behavior may change"):
         ddf2 = dd.read_parquet(
             str(tmpdir),
             engine=read_engine,
@@ -3157,16 +3166,15 @@ def test_split_adaptive_blocksize(tmpdir, blocksize, engine, metadata):
         assert "_metadata" not in files
         path = os.path.join(dirname, "*.parquet")
 
-    with pytest.warns(FutureWarning, match="argument will be deprecated"):
-        ddf2 = dd.read_parquet(
-            path,
-            engine=engine,
-            blocksize=blocksize,
-            split_row_groups="adaptive",
-            calculate_divisions=True,
-            index="index",
-            aggregate_files=True,
-        )
+    ddf2 = dd.read_parquet(
+        path,
+        engine=engine,
+        blocksize=blocksize,
+        split_row_groups="adaptive",
+        calculate_divisions=True,
+        index="index",
+        aggregate_files=True,
+    )
 
     assert_eq(ddf1, ddf2, check_divisions=False)
 
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index cd8bca788..de43cc14a 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import pytest
 
 import dask
+from dask.base import collections_to_dsk
 from dask.core import get_deps
 from dask.order import diagnostics, ndependencies, order
 from dask.utils_test import add, inc
@@ -109,22 +110,19 @@ def test_base_of_reduce_preferred(abcde):
     assert o[(b, 1)] <= 3
 
 
-@pytest.mark.xfail(reason="Can't please 'em all")
 def test_avoid_upwards_branching(abcde):
     r"""
-         a1
-         |
-         a2
-         |
-         a3    d1
-        /  \  /
-      b1    c1
-      |     |
-      b2    c2
-            |
-            c3
-
-    Prefer b1 over c1 because it won't stick around waiting for d1 to complete
+       a1
+       |
+       a2
+       |
+       a3    d1
+      /  \  /
+    b1    c1
+    |     |
+    b2    c2
+          |
+          c3
     """
     a, b, c, d, e = abcde
     dsk = {
@@ -135,11 +133,13 @@ def test_avoid_upwards_branching(abcde):
         (c, 1): (f, (c, 2)),
         (c, 2): (f, (c, 3)),
         (d, 1): (f, (c, 1)),
+        (c, 3): 1,
+        (b, 2): 1,
     }
 
     o = order(dsk)
 
-    assert o[(b, 1)] < o[(c, 1)]
+    assert o[(d, 1)] < o[(b, 1)]
 
 
 def test_avoid_upwards_branching_complex(abcde):
@@ -272,27 +272,24 @@ def test_type_comparisions_ok(abcde):
     order(dsk)  # this doesn't err
 
 
-def test_prefer_short_dependents(abcde):
+def test_favor_longest_critical_path(abcde):
     r"""
 
-         a
-         |
-      d  b  e
-       \ | /
-         c
+       a
+       |
+    d  b  e
+     \ | /
+       c
 
-    Prefer to finish d and e before starting b.  That way c can be released
-    during the long computations.
     """
     a, b, c, d, e = abcde
     dsk = {c: (f,), d: (f, c), e: (f, c), b: (f, c), a: (f, b)}
 
     o = order(dsk)
-    assert o[d] < o[b]
-    assert o[e] < o[b]
+    assert o[d] > o[b]
+    assert o[e] > o[b]
 
 
-@pytest.mark.xfail(reason="This is challenging to do precisely")
 def test_run_smaller_sections(abcde):
     r"""
             aa
@@ -301,36 +298,26 @@ def test_run_smaller_sections(abcde):
      / \ /|  | /
     a   c e  cc
 
-    Prefer to run acb first because then we can get that out of the way
     """
     a, b, c, d, e = abcde
     aa, bb, cc, dd = (x * 2 for x in [a, b, c, d])
 
-    expected = [a, c, b, e, d, cc, bb, aa, dd]
-
-    log = []
-
-    def f(x):
-        def _(*args):
-            log.append(x)
-
-        return _
-
     dsk = {
-        a: (f(a),),
-        c: (f(c),),
-        e: (f(e),),
-        cc: (f(cc),),
-        b: (f(b), a, c),
-        d: (f(d), c, e),
-        bb: (f(bb), cc),
-        aa: (f(aa), d, bb),
-        dd: (f(dd), cc),
+        a: (f,),
+        c: (f,),
+        e: (f,),
+        cc: (f,),
+        b: (f, a, c),
+        d: (f, c, e),
+        bb: (f, cc),
+        aa: (f, d, bb),
+        dd: (f, cc),
     }
-
-    dask.get(dsk, [aa, b, dd])  # trigger computation
-
-    assert log == expected
+    o = order(dsk)
+    assert max(diagnostics(dsk)[1]) <= 4  # optimum is 3
+    # This is a mildly ambiguous example
+    # https://github.com/dask/dask/pull/10535/files#r1337528255
+    assert (o[aa] < o[a] and o[dd] < o[a]) or (o[b] < o[e] and o[b] < o[cc])
 
 
 def test_local_parents_of_reduction(abcde):
@@ -591,12 +578,18 @@ def test_dont_run_all_dependents_too_early(abcde):
     """From https://github.com/dask/dask-ml/issues/206#issuecomment-395873372"""
     a, b, c, d, e = abcde
     depth = 10
-    dsk = {(a, 0): 0, (b, 0): 1, (c, 0): 2, (d, 0): (f, (a, 0), (b, 0), (c, 0))}
+    dsk = {
+        (a, 0): (f, 0),
+        (b, 0): (f, 1),
+        (c, 0): (f, 2),
+        (d, 0): (f, (a, 0), (b, 0), (c, 0)),
+    }
     for i in range(1, depth):
         dsk[(b, i)] = (f, (b, 0))
         dsk[(c, i)] = (f, (c, 0))
         dsk[(d, i)] = (f, (d, i - 1), (b, i), (c, i))
     o = order(dsk)
+
     expected = [3, 6, 9, 12, 15, 18, 21, 24, 27, 30]
     actual = sorted(v for (letter, num), v in o.items() if letter == d)
     assert expected == actual
@@ -723,13 +716,17 @@ def test_order_with_equal_dependents(abcde):
     This DAG has enough structure to exercise more parts of `order`
 
     """
+    # Lower pressure is better but this is where we are right now. Important is
+    # that no variation below should be worse since all variations below should
+    # reduce to the same graph when optimized/fused.
+    max_pressure = 11
     a, b, c, d, e = abcde
     dsk = {}
     abc = [a, b, c, d]
     for x in abc:
         dsk.update(
             {
-                (x, 0): 0,
+                (x, 0): (f, 0),
                 (x, 1): (f, (x, 0)),
                 (x, 2, 0): (f, (x, 0)),
                 (x, 2, 1): (f, (x, 1)),
@@ -753,7 +750,9 @@ def test_order_with_equal_dependents(abcde):
             val = o[(x, 6, i, 1)] - o[(x, 6, i, 0)]
             assert val > 0  # ideally, val == 2
             total += val
-    assert total <= 110  # ideally, this should be 2 * 16 = 32
+    assert total <= 56  # ideally, this should be 2 * 16 == 32
+    pressure = diagnostics(dsk, o=o)[1]
+    assert max(pressure) <= max_pressure
 
     # Add one to the end of the nine bundles
     dsk2 = dict(dsk)
@@ -764,10 +763,12 @@ def test_order_with_equal_dependents(abcde):
     total = 0
     for x in abc:
         for i in range(len(abc)):
-            val = o[(x, 7, i, 0)] - o[(x, 6, i, 1)]
+            val = o[(x, 6, i, 1)] - o[(x, 7, i, 0)]
             assert val > 0  # ideally, val == 3
             total += val
-    assert total <= 138  # ideally, this should be 3 * 16 == 48
+    assert total <= 75  # ideally, this should be 3 * 16 == 48
+    pressure = diagnostics(dsk2, o=o)[1]
+    assert max(pressure) <= max_pressure
 
     # Remove one from each of the nine bundles
     dsk3 = dict(dsk)
@@ -778,10 +779,12 @@ def test_order_with_equal_dependents(abcde):
     total = 0
     for x in abc:
         for i in range(len(abc)):
-            val = o[(x, 6, i, 0)] - o[(x, 5, i, 1)]
-            assert val > 0  # ideally, val == 2
+            val = o[(x, 5, i, 1)] - o[(x, 6, i, 0)]
+            assert val > 0
             total += val
-    assert total <= 98  # ideally, this should be 2 * 16 == 32
+    assert total <= 45  # ideally, this should be 2 * 16 == 32
+    pressure = diagnostics(dsk3, o=o)[1]
+    assert max(pressure) <= max_pressure
 
     # Remove another one from each of the nine bundles
     dsk4 = dict(dsk3)
@@ -789,10 +792,11 @@ def test_order_with_equal_dependents(abcde):
         for i in range(len(abc)):
             del dsk4[(x, 6, i, 0)]
     o = order(dsk4)
-    total = 0
+    pressure = diagnostics(dsk4, o=o)[1]
+    assert max(pressure) <= max_pressure
     for x in abc:
         for i in range(len(abc)):
-            assert o[(x, 5, i, 1)] - o[(x, 5, i, 0)] == 1
+            assert abs(o[(x, 5, i, 1)] - o[(x, 5, i, 0)]) <= 10
 
 
 def test_terminal_node_backtrack():
@@ -986,7 +990,7 @@ def test_eager_to_compute_dependent_to_free_parent():
 
 def test_diagnostics(abcde):
     r"""
-        a1  b1  c2  d1  e1
+        a1  b1  c1  d1  e1
         /|\ /|\ /|\ /|  /
        / | X | X | X | /
       /  |/ \|/ \|/ \|/
@@ -1005,6 +1009,10 @@ def test_diagnostics(abcde):
         (d, 1): (f, (d, 0), (e, 0)),
         (e, 1): (f, (e, 0)),
     }
+    o = order(dsk)
+    assert o[(e, 1)] == len(dsk) - 1
+    assert o[(d, 1)] == len(dsk) - 2
+    assert o[(c, 1)] == len(dsk) - 3
     info, memory_over_time = diagnostics(dsk)
     assert memory_over_time == [0, 1, 2, 3, 2, 3, 2, 3, 2, 1]
     assert {key: val.order for key, val in info.items()} == {
@@ -1067,3 +1075,596 @@ def test_diagnostics(abcde):
         (d, 1): 2,
         (e, 1): 1,
     }
+
+
+def test_xarray_like_reduction():
+    a, b, c, d, e = list("abcde")
+
+    dsk = {}
+    for ix in range(3):
+        part = {
+            # Part1
+            (a, 0, ix): (f,),
+            (a, 1, ix): (f,),
+            (b, 0, ix): (f, (a, 0, ix)),
+            (b, 1, ix): (f, (a, 0, ix), (a, 1, ix)),
+            (b, 2, ix): (f, (a, 1, ix)),
+            (c, 0, ix): (f, (b, 0, ix)),
+            (c, 1, ix): (f, (b, 1, ix)),
+            (c, 2, ix): (f, (b, 2, ix)),
+        }
+        dsk.update(part)
+    for ix in range(3):
+        dsk.update(
+            {
+                (d, ix): (f, (c, ix, 0), (c, ix, 1), (c, ix, 2)),
+            }
+        )
+    o = order(dsk)
+    _, pressure = diagnostics(dsk, o=o)
+    assert max(pressure) <= 9
+
+
+@pytest.mark.parametrize(
+    "optimize",
+    [True, False],
+)
+def test_array_vs_dataframe(optimize):
+    xr = pytest.importorskip("xarray")
+
+    import dask.array as da
+
+    size = 5000
+    ds = xr.Dataset(
+        dict(
+            anom_u=(
+                ["time", "face", "j", "i"],
+                da.random.random((size, 1, 987, 1920), chunks=(10, 1, -1, -1)),
+            ),
+            anom_v=(
+                ["time", "face", "j", "i"],
+                da.random.random((size, 1, 987, 1920), chunks=(10, 1, -1, -1)),
+            ),
+        )
+    )
+
+    quad = ds**2
+    quad["uv"] = ds.anom_u * ds.anom_v
+    mean = quad.mean("time")
+    diag_array = diagnostics(collections_to_dsk([mean], optimize_graph=optimize))
+    diag_df = diagnostics(
+        collections_to_dsk([mean.to_dask_dataframe()], optimize_graph=optimize)
+    )
+    assert max(diag_df[1]) == max(diag_array[1])
+    assert max(diag_array[1]) < 50
+
+
+def test_anom_mean():
+    np = pytest.importorskip("numpy")
+    xr = pytest.importorskip("xarray")
+
+    import dask.array as da
+    from dask.utils import parse_bytes
+
+    data = da.random.random(
+        (260, 1310720),
+        chunks=(1, parse_bytes("10MiB") // 8),
+    )
+
+    ngroups = data.shape[0] // 4
+    arr = xr.DataArray(
+        data,
+        dims=["time", "x"],
+        coords={"day": ("time", np.arange(data.shape[0]) % ngroups)},
+    )
+    data = da.random.random((5, 1), chunks=(1, 1))
+
+    arr = xr.DataArray(
+        data,
+        dims=["time", "x"],
+        coords={"day": ("time", np.arange(5) % 2)},
+    )
+
+    clim = arr.groupby("day").mean(dim="time")
+    anom = arr.groupby("day") - clim
+    anom_mean = anom.mean(dim="time")
+    _, pressure = diagnostics(anom_mean.__dask_graph__())
+    assert max(pressure) <= 9
+
+
+def test_anom_mean_raw():
+    dsk = {
+        ("d", 0, 0): (f, ("a", 0, 0), ("b1", 0, 0)),
+        ("d", 1, 0): (f, ("a", 1, 0), ("b1", 1, 0)),
+        ("d", 2, 0): (f, ("a", 2, 0), ("b1", 2, 0)),
+        ("d", 3, 0): (f, ("a", 3, 0), ("b1", 3, 0)),
+        ("d", 4, 0): (f, ("a", 4, 0), ("b1", 4, 0)),
+        ("a", 0, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 1, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 2, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 3, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 4, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("e", 0, 0): (f, ("g1", 0)),
+        ("e", 1, 0): (f, ("g3", 0)),
+        ("b0", 0, 0): (f, ("a", 0, 0)),
+        ("b0", 1, 0): (f, ("a", 2, 0)),
+        ("b0", 2, 0): (f, ("a", 4, 0)),
+        ("c0", 0, 0): (f, ("b0", 0, 0)),
+        ("c0", 1, 0): (f, ("b0", 1, 0)),
+        ("c0", 2, 0): (f, ("b0", 2, 0)),
+        ("g1", 0): (f, [("c0", 0, 0), ("c0", 1, 0), ("c0", 2, 0)]),
+        ("b2", 0, 0): (f, ("a", 1, 0)),
+        ("b2", 1, 0): (f, ("a", 3, 0)),
+        ("c1", 0, 0): (f, ("b2", 0, 0)),
+        ("c1", 1, 0): (f, ("b2", 1, 0)),
+        ("g3", 0): (f, [("c1", 0, 0), ("c1", 1, 0)]),
+        ("b1", 0, 0): (f, ("e", 0, 0)),
+        ("b1", 1, 0): (f, ("e", 1, 0)),
+        ("b1", 2, 0): (f, ("e", 0, 0)),
+        ("b1", 3, 0): (f, ("e", 1, 0)),
+        ("b1", 4, 0): (f, ("e", 0, 0)),
+        ("c2", 0, 0): (f, ("d", 0, 0)),
+        ("c2", 1, 0): (f, ("d", 1, 0)),
+        ("c2", 2, 0): (f, ("d", 2, 0)),
+        ("c2", 3, 0): (f, ("d", 3, 0)),
+        ("c2", 4, 0): (f, ("d", 4, 0)),
+        ("f", 0, 0): (f, [("c2", 0, 0), ("c2", 1, 0), ("c2", 2, 0), ("c2", 3, 0)]),
+        ("f", 1, 0): (f, [("c2", 4, 0)]),
+        ("g2", 0): (f, [("f", 0, 0), ("f", 1, 0)]),
+    }
+
+    o = order(dsk)
+
+    # The left hand computation branch should complete before we start loading
+    # more data
+    nodes_to_finish_before_loading_more_data = [
+        ("f", 1, 0),
+        ("d", 0, 0),
+        ("d", 2, 0),
+        ("d", 4, 0),
+    ]
+    for n in nodes_to_finish_before_loading_more_data:
+        assert o[n] < o[("a", 1, 0)]
+        assert o[n] < o[("a", 3, 0)]
+
+
+def test_flaky_array_reduction():
+    first = {
+        ("mean_agg-aggregate-10d721567ef5a0d6a0e1afae8a87c066", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 0, 0, 0, 0),
+                ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 0, 0, 0, 0),
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 1, 0, 0, 0),
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 2, 0, 0, 0),
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 0, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 0, 0, 0, 0),
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 0, 0, 0, 0),
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 1, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 1, 0, 0, 0),
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 1, 0, 0, 0),
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 2, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 2, 0, 0, 0),
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 2, 0, 0, 0),
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 3, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 3, 0, 0, 0),
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 3, 0, 0, 0),
+        ),
+        ("mean_agg-aggregate-fdb340546b01334890192fcfa55fa0d9", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 0, 0, 0, 0),
+                ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 0, 0, 0, 0),
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 1, 0, 0, 0),
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 2, 0, 0, 0),
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-23adb4747560e6e33afd63c5bb179709",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 0, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 0, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 1, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 1, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 2, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 2, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 3, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 3, 0, 0, 0),
+            2,
+        ),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 0, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 1, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 2, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 3, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 4, 0, 0, 0): (f, 1),
+        ("mean_agg-aggregate-cc19342c8116d616fc6573f5d20b5762", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 0, 0, 0, 0),
+                ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 0, 0, 0, 0),
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 1, 0, 0, 0),
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 2, 0, 0, 0),
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-0c98c5a4517f58f8268985e7464daace",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 0, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 0, 0, 0, 0),
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 1, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 1, 0, 0, 0),
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 2, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 2, 0, 0, 0),
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 3, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 3, 0, 0, 0),
+        ),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 0, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 1, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 2, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 3, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 4, 0, 0, 0): (f, 1),
+        (
+            "mean_chunk-mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [
+                (
+                    f,
+                    ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 4, 0, 0, 0),
+                    ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 4, 0, 0, 0),
+                )
+            ],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-0c98c5a4517f58f8268985e7464daace",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 4, 0, 0, 0), 2)],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-23adb4747560e6e33afd63c5bb179709",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 4, 0, 0, 0), 2)],
+        ),
+    }
+
+    other = {
+        ("mean_agg-aggregate-e79dd3b9757c9fb2ad7ade96f3f6c814", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 0, 0, 0, 0),
+                ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 0, 0, 0, 0),
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 1, 0, 0, 0),
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 2, 0, 0, 0),
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 0, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 0, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 0, 0, 0, 0),
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 1, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 1, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 1, 0, 0, 0),
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 2, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 2, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 2, 0, 0, 0),
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 3, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 3, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 3, 0, 0, 0),
+        ),
+        ("mean_agg-aggregate-c7647920facf0e557f947b7a6626b7be", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 0, 0, 0, 0),
+                ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 0, 0, 0, 0),
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 1, 0, 0, 0),
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 2, 0, 0, 0),
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 0, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 0, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 1, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 1, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 2, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 2, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 3, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 3, 0, 0, 0),
+            2,
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 0, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 1, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 2, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 3, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 4, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("mean_agg-aggregate-05071ebaabb68a64c180f6f443c5c8f4", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 0, 0, 0, 0),
+                ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 0, 0, 0, 0),
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 1, 0, 0, 0),
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 2, 0, 0, 0),
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 0, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 0, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 1, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 1, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 2, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 2, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 3, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 3, 0, 0, 0),
+            2,
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 0, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 1, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 2, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 3, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 4, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 4, 0, 0, 0), 2)],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [
+                (
+                    f,
+                    ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 4, 0, 0, 0),
+                    ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 4, 0, 0, 0),
+                )
+            ],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 4, 0, 0, 0), 2)],
+        ),
+    }
+    first_pressure = max(diagnostics(first)[1])
+    second_pressure = max(diagnostics(other)[1])
+    assert first_pressure == second_pressure
+
+
+def test_flox_reduction():
+    # TODO: It would be nice to scramble keys to ensure we're not comparing keys
+    dsk = {
+        "A0": (f, 1),
+        ("A1", 0): (f, 1),
+        ("A1", 1): (f, 1),
+        ("A2", 0): (f, 1),
+        ("A2", 1): (f, 1),
+        ("B1", 0): (f, [(f, ("A2", 0))]),
+        ("B1", 1): (f, [(f, ("A2", 1))]),
+        ("B2", 1): (f, [(f, ("A1", 1))]),
+        ("B2", 0): (f, [(f, ("A1", 0))]),
+        ("B11", 0): ("B1", 0),
+        ("B11", 1): ("B1", 1),
+        ("B22", 0): ("B2", 0),
+        ("B22", 1): ("B2", 1),
+        ("C1", 0): (f, ("B22", 0)),
+        ("C1", 1): (f, ("B22", 1)),
+        ("C2", 0): (f, ("B11", 0)),
+        ("C2", 1): (f, ("B11", 1)),
+        ("E", 1): (f, [(f, ("A1", 1), ("A2", 1), ("C1", 1), ("C2", 1))]),
+        ("E", 0): (f, [(f, ("A1", 0), ("A2", 0), ("C1", 0), ("C2", 0))]),
+        ("EE", 0): ("E", 0),
+        ("EE", 1): ("E", 1),
+        ("F1", 0): (f, "A0", ("B11", 0)),
+        ("F1", 1): (f, "A0", ("B22", 0)),
+        ("F1", 2): (f, "A0", ("EE", 0)),
+        ("F2", 0): (f, "A0", ("B11", 1)),
+        ("F2", 1): (f, "A0", ("B22", 1)),
+        ("F2", 2): (f, "A0", ("EE", 1)),
+    }
+    o = order(dsk)
+    assert max(o[("F1", ix)] for ix in range(3)) < min(o[("F2", ix)] for ix in range(3))