Repository: dask/dask
File: dask/order.py (Python)

Code before:

    # Pick strategy
    # Note: We're trying to be smart here by picking a strategy on how to
    # determine the critical path. This is not always clear and we may want to
    # consider just calculating both orderings and picking the one with less
    # pressure. The only concern to this would be performance but at time of
    # writing, the most expensive part of ordering is the prep work (mostly
    # connected roots + sort_key) which can be reused for multiple orderings.

    # Only build this mapping if it's needed. This can be expensive for
    # all-to-all communications, but we normally don't need it in this
    # case
    occurences_grouped_sorted: dict[int, list[Key]] = {}

    def _built_occurrences_group_sorted() -> None:
        if occurences_grouped_sorted:
            return
        occurrences: defaultdict[Key, int] = defaultdict(int)
        for t in leaf_nodes:
            for r in roots_connected[t]:
                occurrences[r] += 1
        occurences_grouped = defaultdict(set)
        for root, occ in occurrences.items():
            occurences_grouped[occ].add(root)
        for k, v in occurences_grouped.items():
            occurences_grouped_sorted[k] = sorted(v, key=sort_key, reverse=True)

    def pick_seed() -> Key | None:
        if not occurences_grouped_sorted:
            _built_occurrences_group_sorted()

        while occurences_grouped_sorted:
            key = max(occurences_grouped_sorted)
            picked_root = occurences_grouped_sorted[key][-1]
            if picked_root in result:
                occurences_grouped_sorted[key].pop()
                if not occurences_grouped_sorted[key]:
                    del occurences_grouped_sorted[key]
                continue
            return picked_root
        return None

    def get_target(longest_path: bool = False) -> Key:
        # Some topologies benefit if the node with the most dependencies
        # is used as first choice, others benefit from the opposite.
        candidates = leaf_nodes
        skey: Callable = sort_key

        preferred_candidates = set()
        if runnable_hull:
            skey = lambda k: (num_needed[k], sort_key(k))

Review comment (performance):
Unrelated to your changes but given performance here has been an issue...

We should be able to use the knowledge that all `v` are root nodes and perform the sort on all roots once.
With this knowledge we should be able to do sort every iteration in N instead of N logN


Code after:
                                            add_to_result(k)
                        else:
                            if (
                                len(dependencies[current]) > 1  # type: ignore
                                and num_needed[current] <= 1
                            ):
                                for k in path:
                                    add_to_result(k)
                            else:
                                known_runnable_paths[current] = [path]
                                runnable_hull.discard(current)
                        break

    # Pick strategy
    # Note: We're trying to be smart here by picking a strategy on how to
    # determine the critical path. This is not always clear and we may want to
    # consider just calculating both orderings and picking the one with less
    # pressure. The only concern to this would be performance but at time of
    # writing, the most expensive part of ordering is the prep work (mostly
    # connected roots + sort_key) which can be reused for multiple orderings.

    def get_target(longest_path: bool = False) -> Key:
        # Some topologies benefit if the node with the most dependencies
        # is used as first choice, others benefit from the opposite.
        candidates = leaf_nodes
        skey: Callable = sort_key
        if reachable_hull:
            skey = lambda k: (num_needed[k], sort_key(k))

        all_leafs_accessible = len(leafs_connected_to_loaded_roots) >= len(candidates)
        if reachable_hull and not all_leafs_accessible:
            # Avoid this branch if we can already access all leafs. Just pick
            # one of them without the expensive selection process in this case.

            candidates = reachable_hull & candidates
            if not candidates:
                # We can't reach a leaf node directly, but we still have nodes
                # with results in memory, these notes can inform our path towards
                # a new preferred leaf node.
                for r in processed_roots:
                    leafs_connected_to_loaded_roots.update(leafs_connected[r])
                processed_roots.clear()

                leafs_connected_to_loaded_roots.intersection_update(leaf_nodes)
                candidates = leafs_connected_to_loaded_roots
            else:
                leafs_connected_to_loaded_roots.update(candidates)

        elif not reachable_hull:
            # We reach a new and independent branch, so throw away previous branch
            leafs_connected_to_loaded_roots.clear()