rjgoyln opened a new pull request, #70620: URL: https://github.com/apache/airflow/pull/70620
## Summary Resolving mapped task kwargs previously issued one database query per `XComArg`, causing N+1 queries for expanded kwargs (especially with nested `zip()`/`concat()`). This change replaces those repeated lookups with a bulk prefetch, addressing the long-standing TODO in `SchedulerDictOfListsExpandInput._get_map_lengths`. Upstream map lengths are now resolved in a fixed number of queries, improving both scheduler expansion and execution API argument binding. ## Changes - Add `prefetch_map_lengths()` to resolve all referenced upstream map lengths in bulk. - Add an optional `lengths` parameter to `get_task_map_length()` and use prefetched results when available. - Prefetch once in `SchedulerDictOfListsExpandInput` and `SchedulerListOfDictsExpandInput`. - Add tests covering query counts and behavior for mapped, unmapped, zipped, and unfinished upstreams. ## Testing - Added query-count tests to verify the optimization. - Existing scheduler and mapping tests pass. - `mypy-airflow-core` and pre-commit checks pass. * closes: #70548 * related: #70548 --> --- ##### Was generative AI tooling used to co-author this PR? <!-- If generative AI tooling has been used in the process of authoring this PR, please change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by". --> - [ ] Yes (please specify the tool below) <!-- Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) --> --- * Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged. * For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed. * When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x). * For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number. -- 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]
