Optimize TaskGroup.topological_sort for reverse-declared Dags

[shahar1]

Human Summary

Further optimization of TaskGroup.topological_sort to handle reverse-declared Dags (where many children are declared before their dependencies) efficiently. This is the follow-up to PR #67288.

Explanation

The idea is to optimize TaskGroup.topological_sort for Dags whose tasks are declared in reverse of their dependency order, for example:

@dag
def dag():
    task3 = EmptyOperator()
    task2 = EmptyOperator()
    task1 = EmptyOperator()

    task1 >> task2 >> task3

Note that we insert task3, then task2, then task1, but their topological order is the other way round.

The issue with this shape is that the current sweep algorithm revisits blocked nodes. In a reversed-declared chain, we end up in a situation where total of N + (N-1) + (N-2) + ... + 1 nodes are revisited, leading to a time complexity of O(N^2).
To solve this, we first use Kahn's algorithm to compute the pass number at which each node would be emitted by the existing sweep algorithm. Once the pass numbers are known, we sort nodes by (pass_number, insertion_order), which produces the exact same ordering as the sweep algorithm while avoiding the repeated rescans.

Since the pass-numbering approach has a higher constant overhead than the sweep algorithm, we only enable it for TaskGroups with a significant number of back edges (nodes_with_back_edge >= 32), where the cost of repeated sweeps is expected to outweigh that overhead (backed up by a benchmark script).

Was generative AI tooling used to co-author this PR?

• Yes — Claude Code following the guidelines

AI Summary

Click here

Addresses the O(N²) worst-case behavior of the greedy-sweep approach on adversarial Dag shapes such as reverse-insertion chains.

Uses a hybrid strategy:

• Forward-declared Dags (common case): greedy multi-pass sweep, O(V + E)
• Reverse-declared Dags (many back edges): pass-number traversal via Kahn's algorithm, O((V + E) log V)
• Mixed Dags where independent children dilute the ratio: the dispatcher also flips once the group has at least 32 back-edge nodes, which catches the padded reverse-chain case raised in review

Both approaches emit the same order: level-by-legacy-pass, ties broken by insertion order.

Benchmark Results

Run the benchmark with: uv run --project task-sdk python dev/bench_topological_sort_comparison.py

See the gist for the benchmark script.

Reverse-Chain Speedup (Worst Case)

N	Sweep-only (ms)	Hybrid (ms)	Speedup
100	0.29	0.07	4.35x
500	6.01	0.43	13.88x
1000	24.15	0.75	32.06x
2000	96.79	1.77	54.58x

Padded Reverse-Chain (Review Case)

Shape	Sweep-only (ms)	Hybrid (ms)	Speedup
1000 reverse + 1000 independent	24.35	1.18	20.72x

Performance Progression

• Before PR Speed up TaskGroup.topological_sort with int-indexed projected sweep #67288 (old Kahn's): ~894ms for reverse-chain N=2000
• After PR Speed up TaskGroup.topological_sort with int-indexed projected sweep #67288 (sweep-only): ~119ms
• With this follow-up (hybrid): ~1.77ms ✅

The dispatcher now switches when a group is clearly back-heavy either by ratio or by an absolute back-edge count, so padded reverse-declared Dags no longer fall back to the quadratic sweep.

Test Plan

• Reran the benchmark from the gist after rebasing onto main
• Added a reverse-declared order-equivalence regression that asserts _sort_via_pass_numbering matches _sweep_projection
• Added padded reverse-chain dispatch regressions for both TaskGroup and SerializedTaskGroup
• uv run --project task-sdk pytest task-sdk/tests/task_sdk/definitions/test_taskgroup.py -k 'reverse_declared_order_matches_sweep or padded_reverse_chain_uses_pass_numbering or topological_sort_shape_correctness' -xvs
• uv run --project airflow-core pytest airflow-core/tests/unit/utils/test_task_group.py -k 'serialized_padded_reverse_chain_uses_pass_numbering or topological_sort_serialized_layered' -xvs

[kaxil]

cc @uranusjr can you review it too plz

[shahar1]

Merging, edited the PR's description with a concise explanation.