Speed up clearing downstream tasks on large Dags

[ColtenOuO]

Summary

partial_subset walks the downstream relatives of every matched task, and for each one
asks whether that relative is itself among the matched tasks:

matched_tasks = [t for t in self.tasks if t.task_id in task_ids]   # a list

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:   # linear scan, inside a nested loop

matched_tasks is a list because SerializedBaseOperator sets __hash__ = None, so the
operators cannot go in a set at all. Its __eq__ returns NotImplemented, which makes
in fall back to an identity scan — cheap per comparison, but still O(len(matched_tasks))
per relative.

That puts the check at O(matched x relatives). Both factors grow with the Dag, so on a Dag
whose tasks mostly reach one another the whole call goes cubic in the task count.

Task ids are unique within a Dag, and the relatives get_flat_relatives yields are the very
objects matched_tasks holds, so the identity scan and a set lookup on task_id answer the
same question. The set lookup is constant time, dropping the check to O(matched+ relatives).

Although it seems like a small change, I looked into it further and gathered some data. Here's what I found:

What drives the cost

The improvement tracks the number of downstream relatives, which is set by the Dag's
depth, not its task count. Four shapes, each measured at the same three task counts.

Chain

Every task reaches every later one, so relatives grow as $\frac{N(N-1)}{2}$. This is the ceiling.

tasks	relatives	before	after	speedup
100	4,950	0.028s	0.009s	3.1x
200	19,900	0.175s	0.027s	6.5x
400	79,800	1.335s	0.096s	13.9x

Doubling the task count multiplies the old timing by ~7.6 and the new one by ~3.6 —
cubic against quadratic.

---

Layered

Width does not help on its own. Connecting each layer fully to the next still lets every
task reach everything downstream of it, so relatives stay near $\frac{N^2}{2}$ and the timings
land almost on top of the chain.

tasks	relatives	before	after	speedup
100	4,000	0.024s	0.008s	3.0x
200	18,000	0.170s	0.027s	6.3x
400	76,000	1.312s	0.116s	11.3x

---

Parallel chains

Splitting one Dag into $K$ independent pipelines divides the relatives by $K$ — several
unrelated flows sharing a Dag file is a common shape, and it still gains meaningfully.

tasks	relatives	before	after	speedup
100	450	0.005s	0.003s	1.7x
200	1,900	0.024s	0.008s	3.0x
400	7,800	0.147s	0.022s	6.7x

---

Fan-out

The control. One task has relatives and the rest have none, so there is almost no
scanning to remove and the change should do nothing — which is what it does.

tasks	relatives	before	after	speedup
100	99	0.003s	0.003s	1.0x
200	199	0.007s	0.006s	1.2x
400	399	0.022s	0.012s	1.8x

Reading the four together

The speedup ranks exactly with the relative count — 79,800 → 13.9x, 76,000 → 11.3x,
7,800 → 6.7x, 399 → 1.8x — and within each shape the ratio doubles as the task count
doubles: 3.1x, 6.5x, 13.9x on the chain. That is what one extra linear factor looks like,
and the flat fan-out row is the check that the gain is coming from the scan rather than
from the benchmark.

Benchmark script

from __future__ import annotations

import time

import pendulum

from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk import DAG
from airflow.serialization.serialized_objects import DagSerialization


def _serialize(dag):
    return DagSerialization.from_dict(DagSerialization.to_dict(dag))


def build_chain(n: int):
    """t0 -> t1 -> ... -> tn. Every task reaches every later one."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        previous = None
        for i in range(n):
            task = EmptyOperator(task_id=f"t{i:04d}")
            if previous is not None:
                previous >> task
            previous = task
    return _serialize(dag)


def build_parallel_chains(n: int, streams: int = 10):
    """`streams` independent chains side by side."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        tails: dict[int, object] = {}
        for i in range(n):
            task = EmptyOperator(task_id=f"t{i:04d}")
            stream = i % streams
            if (previous := tails.get(stream)) is not None:
                previous >> task
            tails[stream] = task
    return _serialize(dag)


def build_layered(n: int, width: int = 20):
    """Layers of `width` tasks, each fully connected to the next layer."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        tasks = [EmptyOperator(task_id=f"t{i:04d}") for i in range(n)]
        for start in range(0, n - width, width):
            for upstream in tasks[start : start + width]:
                for downstream in tasks[start + width : start + 2 * width]:
                    upstream >> downstream
    return _serialize(dag)


def build_fanout(n: int):
    """One root feeding every other task -- the shallow extreme."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        root = EmptyOperator(task_id="root")
        for i in range(n - 1):
            root >> EmptyOperator(task_id=f"t{i:04d}")
    return _serialize(dag)


SHAPES = [
    ("chain", build_chain),
    ("parallel chains x10", build_parallel_chains),
    ("layered (w=20)", build_layered),
    ("fan-out", build_fanout),
]


def timed(dag, task_ids) -> float:
    started = time.monotonic()
    dag.partial_subset(task_ids=task_ids, include_downstream=True, include_upstream=False)
    return time.monotonic() - started


print(f"{'shape':<16} {'tasks':>6} {'relatives':>11} {'partial_subset':>16}")
for label, build in SHAPES:
    for n in (100, 200, 400):
        dag = build(n)
        task_ids = {t.task_id for t in dag.tasks}
        relatives = sum(len(t.get_flat_relatives(upstream=False, depth=None)) for t in dag.tasks)
        best = min(timed(dag, task_ids) for _ in range(3))
        print(f"{label:<16} {n:>6} {relatives:>11,} {best:>15.3f}s")
    print()

Result

shape	relatives at N=400	speedup at N=100	at N=200	at N=400
chain	79,800	3.1x	6.5x	13.9x
layered (width 20)	76,000	3.0x	6.3x	11.3x
parallel chains x10	7,800	1.7x	3.0x	6.7x
fan-out	399	1.0x	1.2x	1.8x

[jason810496]

Thanks for the improvement.

Why not make the existing matched_tasks a set (which is still a hash map under the hook with constant time lookup)?

By this way, we can avoid the additional memory usage.

[ColtenOuO]

It looks like we might run into an issue using a set here. Since SerializedBaseOperator has __hash__ = None (and __eq__ returns NotImplemented), the objects themselves are unhashable and can't be stored in a set.

As a workaround, we could take a different approach: rely on matched_task_ids: set[str] as the primary data structure and skip creating the matched_tasks (list of objects) altogether. We'd just look up the object properties only when needed (ex. self.task_dict[tid]).

I've changed the code to this approach. Let me know if this is what you were looking for >< thanks again!