seanmuth opened a new pull request, #72395:
URL: https://github.com/apache/airflow/pull/72395
closes: #72393
## Problem
`DagModelOperation.find_orm_dags` in
`airflow-core/src/airflow/dag_processing/collection.py` eagerly loads five
different one-to-many collections on `DagModel` (`tags`,
`schedule_asset_references`, `schedule_asset_alias_references`,
`task_outlet_asset_references`, `dag_owner_links`) using `joinedload()` in a
single query:
```python
stmt = with_row_locks(
(
select(DagModel)
.options(joinedload(DagModel.tags, innerjoin=False))
.where(DagModel.dag_id.in_(self.dags))
.options(joinedload(DagModel.schedule_asset_references))
.options(joinedload(DagModel.schedule_asset_alias_references))
.options(joinedload(DagModel.task_outlet_asset_references))
.options(joinedload(DagModel.dag_owner_links))
),
of=DagModel,
session=session,
)
```
Combining multiple `joinedload()` calls on one-to-many collections into a
single query produces a cartesian product across those collections. This was
confirmed live on a production deployment: 500 input `dag_id`s produced
**3,907** result rows via `EXPLAIN (ANALYZE, BUFFERS)`. The database itself
executed the query quickly (16ms) — the real cost lands on the client, which
has to receive, deserialize, and de-duplicate (via `.unique()`) every exploded
row, including redundant copies of `DagModel`'s JSON columns
(`partition_mapper_info`, `asset_expression`, `deadline`) once per duplicate.
This was observed causing multi-second processing delays per call in production
for an async SQLAlchemy/asyncpg client, but it affects any client — sync or
async — syncing a DAG set with a meaningful number of tags, asset references,
or owner links.
## Fix
Switch all five relationships to `selectinload()`. Instead of one big join,
this issues one follow-up `WHERE dag_id IN (...)` query per collection — the
same eager-loading outcome (all five collections populated on the returned
`DagModel` objects), but each query returns only rows that actually exist, with
no multiplication.
This is a real tradeoff worth being explicit about: `find_orm_dags` is
called twice per `DAG.bulk_write_to_db()` invocation (once to look up existing
DagModels, once to refetch after flushing new assets), so this goes from 2
statements total (1 per call) to up to 12 (1 base + 5 selectin, times 2 calls).
More round trips, but each one is cheap, returns exactly the right number of
rows, and avoids the client-side deserialize/dedup cost of the exploded joined
result.
Also dropped `innerjoin=False` from the `tags` load — that's a
`joinedload`-specific knob to prevent an inner join from silently dropping DAGs
with no tags. `selectinload` has no equivalent concern since it's a separate
query keyed off the dag_ids already returned by the base select, so a DAG with
zero tags is unaffected either way.
## Testing
- Ran the full `airflow-core/tests/unit/dag_processing/test_collection.py`
and `test_manager.py` suites against both sqlite and Postgres backends — all
passing.
- `test_manager.py` has a pinned per-call SQL statement budget
(`FIXED_PER_CALL`) that tracks exactly this kind of change. Updated it from 9
to 19 with an inline comment explaining the +10 (5 extra `selectin` statements
per `find_orm_dags()` call, times 2 calls per persistence call), so a future
statement-count regression there is easy to diagnose rather than a mystery.
- Local verification: generated a synthetic file with 400
dynamically-generated DAGs (`globals()[dag_id] = dag` pattern), each with 5
tags, 2 task-level asset outlets, and 2 owner links, then ran it through
`SerializedDAG.bulk_write_to_db` (the same path the scheduler's dag processor
uses) against a local Postgres instance.
- Confirmed the unpatched `joinedload` query returns **8,000** rows for
those 400 `dag_id`s (a 20x explosion from the 5×2×2 collection sizes) —
directionally consistent with, and worse than, the 500→3,907 (~7.8x) production
case.
- Confirmed the patched version persists all 400 DAGs correctly — tags,
owner links, and task outlet asset references all present and correct on
spot-checked DAGs — with no explosion.
- Timed steady-state re-sync of the same 400-DAG file 5x each way on the
same local Postgres instance (near-zero network latency, so this understates
the real-world win): unpatched avg ~0.75s, patched avg ~0.54s (~28% faster).
The gap should be substantially larger for any real deployment with actual
client-DB network latency and per-row deserialization overhead, which is
exactly the scenario described in the issue.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
