Speed up TaskGroup.topological_sort with int-indexed projected sweep by shahar1 · Pull Request #67288 · apache/airflow

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

shahar1 · 2026-05-21T15:39:42Z

Summary

Replaces the modified-Kahn body of TaskGroup.topological_sort (and the mirror in SerializedTaskGroup) with a projected sweep:

Each child's per-task upstream_task_ids is projected onto sibling-level integer indices once up front (_project_child_deps).
The sweep (_sweep_projection) then runs against the projection with a bytearray-backed emission flag - per-edge work happens once per sort instead of once per outer-loop pass.

Emission order is identical to the previous implementation. The existing order-sensitive tests (test_topological_sort1/2, test_topological_nested_groups, test_topological_group_dep) cover the contract.

Benchmark

Min of five runs, N=2000 children, ms per call. Source: https://gist.github.com/shahar1/9c61dc9f34f7e77cd29cfb9d67af7ceb

Shape	main	PR	speedup
chain (forward-declared)	1.39	0.75	1.84x
diamond (root → many → sink)	2.20	0.86	2.56x
independent (no deps)	0.65	0.35	1.84x
layered (4-layer full bipartite)	230.96	67.15	3.44x
nested (depth-3, walked across all groups)	17.78	4.77	3.73x
rev-chain (adversarial reverse insertion)	894.04	114.16	7.83x

The 7.8x on rev-chain comes from the projection precompute alone (one projection vs N projections); algorithmic complexity is unchanged and remains O(N²) on that shape until the follow-up PR.

The nested row times a recursive walk across all 39 nested groups (depth=3, ~222 tasks each); the root-only timing was dominated by setup overhead and was previously labelled ~noise in this table.

Test plan

task-sdk/tests/task_sdk/definitions/test_taskgroup.py + test_dag.py — 145 passed (18 new shape-correctness cases)
airflow-core/tests/unit/utils/test_task_group.py — 20 passed (1 new serialization round-trip case)
airflow-core/tests/unit/serialization/ (task-group / topological filter) — 4 passed
prek run mypy-task-sdk — passed
prek run mypy-airflow-core — passed
prek run --from-ref upstream/main --stage pre-commit — passed

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

Yes — Claude Code (Opus 4.7)

Generated-by: Claude Code (Opus 4.7) following the guidelines

Copilot

Pull request overview

This PR optimizes TaskGroup.topological_sort (and the serialized equivalent) by pre-projecting upstream dependencies to sibling indices once per call, then performing a sweep using an efficient emitted-flag structure. This aims to reduce repeated per-edge work during sorting while preserving the existing emission order.

Changes:

Replaced the previous modified-Kahn loop in TaskGroup.topological_sort with a “project then sweep” implementation, including explicit cycle detection.
Mirrored the same algorithm into SerializedTaskGroup.topological_sort for consistency/performance.
Added Task SDK tests that validate topological correctness across multiple DAG shapes, plus a newsfragment entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
task-sdk/tests/task_sdk/definitions/test_taskgroup.py	Adds cross-shape correctness tests for `TaskGroup.topological_sort`.
task-sdk/src/airflow/sdk/definitions/taskgroup.py	Implements projected dependency computation and sweep-based topo sort in Task SDK `TaskGroup`.
airflow-core/src/airflow/serialization/definitions/taskgroup.py	Mirrors the new topo sort algorithm for serialized task groups and adds cycle detection behavior.
airflow-core/newsfragments/67288.improvement.rst	Documents the `TaskGroup.topological_sort` performance improvement.

The previous modified-Kahn implementation re-derived each child's upstream edges (materializing ``upstream_list`` and walking ``parent_group``) on every outer-loop pass. Project per-task upstream IDs onto sibling-level integer indices once up front, then run a greedy multi-pass sweep against that projection with a ``bytearray`` emission flag. Emission order is identical to the previous implementation; existing order-sensitive tests cover the contract. Same change is mirrored in SerializedTaskGroup.

Adds a round-trip test for SerializedTaskGroup.topological_sort (the serialization variant was previously untested), rewrites the newsfragment in user-facing terms, and cleans up a stale reference and type annotation in the task-sdk shape tests.

Per Copilot review feedback, the previous wording said cycles are caught at "deserialization time" — but DAG.check_cycle runs at DAG parse time (via dagbag loading), not during from_dict/from_json. Reword to describe cycles reaching this code path as malformed serialized data, with the defensive ValueError still raised on detection.

Tighten _project_child_deps and the topological_sort body: drop the per-edge child_idx comparison in favour of a single set.discard at the end, inline the get_task/.task_group chain, and replace the pre-alloc loop with a list comprehension. Net -15 lines per file with a small (~10%) speedup on the layered shape and no regression elsewhere. The hot _sweep_projection inner loop is left intact — earlier attempts to extract its duplicated body into a closure cost 10-30% on independent / layered shapes and were reverted.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Drop n=500 from the parametrize grid. The algorithm has no n-dependent branches, so n=100 covers every code path; n=500 only re-runs the same loops with more iterations and added ~0.7s to the file without exercising new correctness behaviour (per Copilot review feedback on PR apache#67288).

kaxil

cc @uranusjr

…-up to apache#67288)

…ected sweep (#67288) (#67687) * Speed up TaskGroup.topological_sort with int-indexed projected sweep (backport #67288) * Remove __getitem__ and NodeNotFound tests not available in v3-2-test These tests depend on TaskGroup.__getitem__ and NodeNotFound exception which don't exist in v3-2-test. Keep topological_sort shape correctness tests which are for the actual optimization. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>

…-up to apache#67288)

…deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite).

…dependencies (#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

…dependencies (#69933) (#70591) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: #65291 Related: #67964 (closed for inactivity, written against the topological_sort implementation before PR #67288/#67688 rewrote it) and #65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> (cherry picked from commit d7aa929) # Conflicts: # airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py Co-authored-by: Hemkumar Chheda <95332229+hkc-8010@users.noreply.github.com>

…dependencies (apache#69933) * Fix grid/graph view topological sort for group-level and cross-group deps TaskGroup._project_child_deps only looked at a group's own upstream_task_ids, which stays empty for a direct group-to-group dependency (list or individual `>>`) and for a task-level dependency that crosses into another group's entry task. Both cases sorted the group as if it had no upstream at all. Now also pulls in the group's upstream_group_ids and its root tasks' upstream task ids before projecting sibling dependencies. Applied to both the serialization-layer sort and the mirrored design-time sort in task-sdk. closes: apache#65291 Related: apache#67964 (closed for inactivity, written against the topological_sort implementation before PR apache#67288/apache#67688 rewrote it) and apache#65639 (draft, same issue, also predates the rewrite). * Address review feedback: cache get_task_group_dict, describe test intent not issue numbers viiccwen pointed out that fetching the group map inside topological_sort() rebuilds the whole DAG's group tree on every nested group's own call, turning a render with G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance (kept behind a small private helper since methodtools.lru_cache has no type stubs and would otherwise widen the public method's return type to Any for every caller). Also reworded test comments/docstrings that cited issue numbers to describe what's actually being verified instead. * Hoist common logic into shared lib * Remove caching on get_task_group_dict * Remove stale get_task_group_dict cache tests The cache these tests asserted was removed in the previous commit, so the identity check and the _get_task_group_dict_cached.cache_info() assertions no longer apply. * Add call-level task group memo to reduce calc * Tidy Typy * Add test for task group memoing --------- Co-authored-by: TP <uranusjr@apache.org> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com>

boring-cyborgBot added area:DAG-processing area:task-sdk labels May 21, 2026

shahar1 force-pushed the taskgroup-topo-sort-projection branch from 90e10cc to 204a0e0Compare May 21, 2026 19:18

shahar1 requested a review from Copilot May 21, 2026 19:18

Copilotstarted reviewing on behalf of shahar1 May 21, 2026 19:18 View session

CopilotAI reviewed May 21, 2026
View reviewed changes

Comment threadtask-sdk/tests/task_sdk/definitions/test_taskgroup.py Outdated
Comment threadairflow-core/src/airflow/serialization/definitions/taskgroup.py Outdated

shahar1 added 6 commits May 21, 2026 23:21

Add newsfragment
77457a4

Add unit tests for correctness
310aa0c

shahar1 force-pushed the taskgroup-topo-sort-projection branch from 62325c1 to 411f90eCompare May 21, 2026 20:22

shahar1 requested a review from Copilot May 21, 2026 20:30

Copilotstarted reviewing on behalf of shahar1 May 21, 2026 20:30 View session

CopilotAI reviewed May 21, 2026
View reviewed changes

Comment threadtask-sdk/tests/task_sdk/definitions/test_taskgroup.py
Comment threadairflow-core/newsfragments/67288.improvement.rst

shahar1 marked this pull request as ready for review May 21, 2026 20:46

shahar1 requested review from amoghrajesh, ashb, bolkedebruin and kaxil as code owners May 21, 2026 20:46

kaxil approved these changes May 22, 2026
View reviewed changes

Merge branch 'main' into taskgroup-topo-sort-projection
96e742b

shahar1 merged commit 7bf71a3 into apache:mainMay 27, 2026
143 checks passed

shahar1 deleted the taskgroup-topo-sort-projection branch May 27, 2026 18:39

This was referenced May 29, 2026
[v3-2-test] Speed up TaskGroup.topological_sort with int-indexed projected sweep (#67288) #67687
Merged
Optimize TaskGroup.topological_sort for reverse-declared Dags #67688
Merged

shahar1 added a commit to shahar1/airflow that referenced this pull request May 29, 2026
Optimize TaskGroup.topological_sort for reverse-declared DAGs (follow…
8506a3f
…-up to apache#67288)

shahar1 added a commit to shahar1/airflow that referenced this pull request May 29, 2026
Optimize TaskGroup.topological_sort for reverse-declared DAGs (follow…
f141482
…-up to apache#67288)

vatsrahul1001 added this to the Airflow 3.3.0 milestone Jun 15, 2026

hkc-8010 mentioned this pull request Jul 15, 2026
Fix grid/graph view topological sort for group-level and cross-group dependencies #69933
Merged
1 task

jason810496 mentioned this pull request Jul 28, 2026
[v3-3-test] Fix grid/graph view topological sort for group-level and cross-group dependencies (#69933) #70591
Merged

vatsrahul1001 mentioned this pull request Aug 7, 2026
Status of testing of Apache Airflow 3.3.1rc2 #71274
Closed