[GitHub] [airflow] uranusjr commented on a diff in pull request #32592: Speed up calculation of leaves and roots for task groups

["uranusjr (via GitHub)" <gi...@apache.org>]

uranusjr commented on code in PR #32592:
URL: https://github.com/apache/airflow/pull/32592#discussion_r1263355739


##########
airflow/utils/task_group.py:
##########
@@ -357,31 +357,35 @@ def get_roots(self) -> Generator[BaseOperator, None, None]:
         Returns a generator of tasks that are root tasks, i.e. those with no upstream
         dependencies within the TaskGroup.
         """
-        for task in self:
-            if not any(self.has_task(parent) for parent in task.get_direct_relatives(upstream=True)):
+        tasks = list(self)
+        ids = {x.task_id for x in tasks}
+        for task in tasks:
+            if not any(parent.task_id in ids for parent in task.get_direct_relatives(upstream=True)):

Review Comment:
   This can use `get_direct_relative_ids` or `upstream_task_ids` to further save some calculation (I think we can even eliminate the `any` call with set operation)



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org