Skip to content

Batch and memoize upstream map-length lookups for mapped stub arg bindings #70548

Description

@jason810496

Background

Since #69757, ti_run derives per-map-index arg bindings for mapped @task.stub tasks via SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes_get_map_lengths (airflow-core/src/airflow/models/expandinput.py). _get_map_lengths carries a pre-existing TODO:

# TODO: This initiates one database call for each XComArg. Would it be
# more efficient to do one single db call and unpack the value here?

Upstream map lengths are immutable once populated for a DagRun, yet every ti_run of every map index re-queries them: a stub expanded over K XComArg kwargs producing N map indexes issues N×K identical get_task_map_length queries per DagRun, on the hot ti_run route. (Stubs expanding over a single kwarg short-circuit and skip the lookups; literal kwargs cost len() only.)

What needs to happen

  1. Batch the per-kwarg length lookups into a single query (addressing the pre-existing _get_map_lengths TODO).
  2. Consider memoizing resolved lengths per (dag_id, run_id, task_id) so repeated ti_run calls for sibling map indexes stop re-querying.

Acceptance criteria

  • One ti_run of a mapped stub issues at most one map-length query.
  • Existing sub-index decomposition tests (airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py) stay green.

Context

Metadata

Metadata

Assignees

Type

No type

Projects

Status
No status

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions