ColtenOuO opened a new pull request, #71203:
URL: https://github.com/apache/airflow/pull/71203
### 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:
```python
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
<img width="1348" height="320" alt="image"
src="https://github.com/user-attachments/assets/9742fcac-e2e2-4c1b-8fea-8419929d4fdd";
/>
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
<img width="769" height="468" alt="image"
src="https://github.com/user-attachments/assets/b8ec10bd-6ad9-40cb-827a-70b6e6b30d6e";
/>
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
<img width="558" height="475" alt="image"
src="https://github.com/user-attachments/assets/50ef67f5-3c4d-49dc-87d2-42ef11535a29";
/>
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
<img width="452" height="511" alt="image"
src="https://github.com/user-attachments/assets/1ed98cc7-0537-49c0-8aa8-7266942c3c71";
/>
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.
<details>
<summary>Benchmark script</summary>
```python
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()
```
</details>
### 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** |
--
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]
