hkc-8010 commented on code in PR #69933: URL: https://github.com/apache/airflow/pull/69933#discussion_r3605630793
########## airflow-core/src/airflow/serialization/definitions/taskgroup.py: ########## Review Comment: Good catch, confirmed this was O(G^2) — task_group_to_dict() recurses into every nested group and calls topological_sort() on each independently, and the Task SDK's own DAG.topological_sort() does the same. Couldn't pass the map through as a literal parameter since topological_sort() is a public zero-arg method invoked externally by several independent recursive callers (task_group_to_dict, dag_edges, DAG.topological_sort()) — changing its signature would ripple into all of them for no real gain. Instead memoized get_task_group_dict() itself with methodtools.lru_cache, keyed per instance: dag.task_group is the same object for every nested group in one DAG, so the first call builds the map and every subsequent nested call is an O(1) hit. Applied the same fix in task-sdk. One wrinkle: methodtools.lru_cache has no type stubs, so decorating get_task_group_dict directly widened its return type to Any for every caller, which actually broke mypy downstream in dag_edges.py. Kept the memoized logic behind a private helper with an explicitly-typed public wrapper so mypy still trusts the declared return type. Added a test asserting the cache is actually hit (not rebuilt) across nested topological_sort() calls, in both packages. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
