GertL opened a new pull request, #74568:
URL: https://github.com/apache/airflow/pull/74568
A fan-out partition mapper can turn one asset event into up to
`partition_mapper_max_downstream_keys` downstream keys per target Dag (default
1000, and a mapper may set its own higher `max_downstream_keys`).
`_queue_partitioned_dags` called `_get_or_create_apdr` once per key, and each
call:
- re-took the `AssetModel` row lock (`SELECT ... FOR KEY SHARE`, or a no-op
`UPDATE` + flush on SQLite),
- ran its own "latest APDR for this key" query,
- flushed after creating or updating the APDR.
That is 2–4 round trips per key, all inside one transaction while the lock
is held. The `AssetModel` row was also re-selected for every target Dag,
although `asset_id` does not change inside the loop.
### Change
- New `AssetManager._get_or_create_apdrs(target_keys=...)` resolves every
key of a target Dag in one pass:
- takes the `AssetModel` lock once (a row lock is held until the
transaction ends, so re-taking it per key added nothing),
- reads the newest APDR of every key in one query (`max(id)` grouped by
`partition_key`), chunked at 500 keys to stay under SQLite's bound-parameter
limit,
- adds new APDRs and flushes once, so their ids are available for the
`PartitionedAssetKeyLog` rows.
- Results are returned in key order, and a key repeated in the list resolves
to the same APDR, matching the previous one-at-a-time behaviour.
- The `partition_date` reconciliation rules (adopt when none, suppress on
conflict, keep when agreeing) are unchanged.
- `_get_or_create_apdr` is kept as a single-key wrapper, so its existing
tests, including the concurrent-creation race test, run unchanged.
- `_queue_partitioned_dags` loads the `AssetModel` once instead of once per
target Dag.
### Effect
Resolving existing keys, measured with `count_queries` against the previous
implementation:
| keys | before | after |
|---|---|---|
| 10 | 20 | 2 |
| 200 | 400 | 2 |
(The previous implementation also flushed once per key.)
### Tests
- `test_get_or_create_apdrs_reuses_pending_and_dedupes_in_order`: key order,
reuse of a pending APDR, repeated keys.
- `test_get_or_create_apdrs_query_count_does_not_grow_with_keys`: resolving
200 existing keys costs the same number of queries as resolving 2.
- Existing `tests/unit/assets/test_manager.py` passes (60 tests), as do the
partition and asset tests in `test_scheduler_job.py`, `test_taskinstance.py`,
`test_partitioned_dag_runs.py` and the partition mapper tests. Run locally on
SQLite only; the Postgres/MySQL row-lock path is left to CI.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
**Enola**: enola provided all the insights
https://github.com/enola-labs/enola
Generated-by: Claude Code following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
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]
