From 890028b958936846e0a7dddcf350df64e7064b4c Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Wed, 5 Aug 2026 19:44:45 +0000 Subject: [PATCH] Speed up clearing downstream tasks on large Dags partial_subset decided whether each downstream relative was already among the matched tasks by scanning a list, once per relative. Both the list and the relative count grow with the Dag, so the check cost O(matched x relatives) -- cubic in the task count for a Dag whose tasks mostly reach one another, which is exactly what clearing with downstream on a deep Dag looks like. Task ids are unique within a Dag, and the relatives are the same objects the matched list already holds, so the identity scan answers the same question as a constant-time lookup against a set of those ids. --- airflow-core/src/airflow/serialization/definitions/dag.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py index 8ed0fee2ccabd..fbd33cb66bf7b 100644 --- a/airflow-core/src/airflow/serialization/definitions/dag.py +++ b/airflow-core/src/airflow/serialization/definitions/dag.py @@ -303,12 +303,13 @@ def is_task(obj) -> TypeIs[SerializedOperator]: else: matched_tasks = [t for t in self.tasks if t.task_id in task_ids] + matched_task_ids = {t.task_id for t in matched_tasks} also_include_ids: set[str] = set() for t in matched_tasks: if include_downstream: for rel in t.get_flat_relatives(upstream=False, depth=depth): also_include_ids.add(rel.task_id) - if rel not in matched_tasks: # if it's in there, we're already processing it + if rel.task_id not in matched_task_ids: # if it's in there, we're already processing it # need to include setups and teardowns for tasks that are in multiple # non-collinear setup/teardown paths if not rel.is_setup and not rel.is_teardown: