Batch and memoize upstream map-length lookups for mapped stub arg bindings
- Dominant language
- Python
- Stars
- 46.9k
- Forks
- 17.8k
- Avg merge
- 2d 9h
- Merged PRs (30d)
- 472
Description
### Background
Since https://github.com/apache/airflow/pull/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:
```python
# 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
- Originating PR: https://github.com/apache/airflow/pull/69757 (review finding)
- Caller: `airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py` (`_resolve_mapped_stub_arg_bindings`)
Contributor guide
Assessment
This issue has not been assessed yet.