Skip to content

Batch upstream map-length lookups when resolving mapped kwargs - #70620

Open
rjgoyln wants to merge 5 commits into
apache:mainfrom
rjgoyln:batch-map-length-lookups
Open

Batch upstream map-length lookups when resolving mapped kwargs#70620
rjgoyln wants to merge 5 commits into
apache:mainfrom
rjgoyln:batch-map-length-lookups

Conversation

@rjgoyln

@rjgoyln rjgoyln commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

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.

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)
    Drafted-by: Claude Code (Opus 5)

  • Read the 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) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@rjgoyln
rjgoyln marked this pull request as ready for review July 29, 2026 05:25
@uranusjr

Copy link
Copy Markdown
Member

This only covers the batching part, not memoizing, so “closing” is wrong.

@rjgoyln

rjgoyln commented Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

This only covers the batching part, not memoizing, so “closing” is wrong.

OK. I got it. Thanks.

Comment thread airflow-core/src/airflow/serialization/definitions/xcom_arg.py Outdated
Comment thread airflow-core/src/airflow/serialization/definitions/xcom_arg.py Outdated
Comment thread airflow-core/src/airflow/serialization/definitions/xcom_arg.py
Comment thread airflow-core/src/airflow/serialization/definitions/xcom_arg.py Outdated
Comment thread airflow-core/src/airflow/serialization/definitions/xcom_arg.py Outdated
rjgoyln added 3 commits July 29, 2026 20:00
The local imports were carried over from the neighbouring lookup rather than
required by anything: no import cycle forces them, and airflow.models is
already a top-level dependency of this module. Feeding update() straight from
the result rows avoids allocating a mapping that is discarded immediately.
@rjgoyln
rjgoyln force-pushed the batch-map-length-lookups branch from 0e9b9c6 to 4eeacca Compare July 29, 2026 12:08
Comment thread airflow-core/src/airflow/serialization/definitions/xcom_arg.py Outdated
@uranusjr

Copy link
Copy Markdown
Member

Is it possible to further reduce the query with UNION? (SQLAlchemy union_all) I feel there’s a chance.

Comment thread airflow-core/tests/unit/serialization/definitions/test_xcom_arg.py
@rjgoyln

rjgoyln commented Jul 31, 2026

Copy link
Copy Markdown
Contributor Author

Is it possible to further reduce the query with UNION? (SQLAlchemy union_all) I feel there’s a chance.

Thanks for the suggestion. I benchmarked the UNION ALL approach.

Advantages

  • Improves performance by about 35% (~0.8 ms) on PostgreSQL 16.

Disadvantages

  • Removes short-circuiting, requiring the database to compute XCom counts for all upstreams and Python to discard unfinished results afterward.
  • Makes the query and unpacking logic significantly more complex.

I have the prototype ready. Let me know if you think the 0.8ms speedup is worth the added complexity, and I can push it!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants