uranusjr commented on code in PR #31921:
URL: https://github.com/apache/airflow/pull/31921#discussion_r1230802049
##########
airflow/models/abstractoperator.py:
##########
@@ -163,30 +163,20 @@ def get_flat_relative_ids(self, *, upstream: bool =
False) -> set[str]:
:param upstream: Whether to look for upstream or downstream relatives.
"""
- dag = self.get_dag()
- if not dag:
- return set()
+ return {x.task_id for x in self.get_flat_relatives(upstream=upstream)}
- relatives: set[str] = set()
-
- task_ids_to_trace = self.get_direct_relative_ids(upstream)
- while task_ids_to_trace:
- task_ids_to_trace_next: set[str] = set()
- for task_id in task_ids_to_trace:
- if task_id in relatives:
- continue
-
task_ids_to_trace_next.update(dag.task_dict[task_id].get_direct_relative_ids(upstream))
- relatives.add(task_id)
- task_ids_to_trace = task_ids_to_trace_next
-
- return relatives
-
- def get_flat_relatives(self, upstream: bool = False) ->
Collection[Operator]:
+ def get_flat_relatives(self, upstream: bool = False) -> set[Operator]:
"""Get a flat list of relatives, either upstream or downstream."""
dag = self.get_dag()
if not dag:
return set()
- return [dag.task_dict[task_id] for task_id in
self.get_flat_relative_ids(upstream=upstream)]
+
+ def get_relatives(task):
+ for rel_task_id in task.get_direct_relative_ids(upstream=upstream):
+ yield (rel_task := dag.task_dict[rel_task_id])
+ yield from get_relatives(rel_task)
Review Comment:
IIRC there were issues when we recursively look up dependencies when the DAG
is large (because this function is called somewhere very deep in the scheduler
stack or something, don’t remember the details). This needs to be iterative.
--
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]
