dstandish opened a new pull request, #69565: URL: https://github.com/apache/airflow/pull/69565
Removes two per-index N+1 query patterns from dynamic task mapping expansion (`TaskMap.expand_mapped_task` and `DagRun._revise_map_indexes_if_mapped`), both of which make a wide fan-out a one-time scheduler stall that scales linearly with map width: - **Batch the inserts.** The per-index `session.merge()` loop that creates the new mapped `TaskInstance` rows issued a merge-load plus a post-merge reload SELECT per index. Replaced with a plain `session.add()` per index and a single `session.flush()` at the end of the loop — the in-tree `# TODO: Make more efficient with bulk_insert_mappings/bulk_save_mappings`. - **Prime `dag_run`.** Later in the same `update_state` cycle, dependency evaluation (`NotPreviouslySkippedDep`) calls `ti.get_dagrun()` per mapped TI. Since `dag_run` isn't loaded on the newly-created instances, each call issued its own `SELECT`. All mapped TIs in a run share one `DagRun`, already in scope during expansion, so it's primed with `set_committed_value` (the same helper `get_dagrun` itself uses to cache) instead of reloading it per TI. ## Test plan - [x] Added a wide XCom-driven expand test (~20 items) asserting correct count, contiguous `map_index` 0..N-1, schedulable state, and that batched instances are usable (have `.task`, are in the session) - [x] Added a mid-run length-grow test for `_revise_map_indexes_if_mapped` - [x] Added a test asserting the `dag_run` relationship is primed after expansion - [x] `prek run --from-ref main --stage pre-commit` passes - [x] `mypy-airflow-core` passes - [x] Existing `test_taskinstance.py` / `test_dagrun.py` mapped-task-expansion tests pass --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes — Claude Code (Sonnet 5) and Claude Opus 4.8 (1M context) Generated-by: Claude Code (Sonnet 5) and Claude Opus 4.8 (1M context) 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]
