apache / apache/airflow

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

Open
#70,548 0 comments 0 reactions 1 assignee Claimed by @jason810496 View on GitHub
area:core area:dynamic-task-mapping area:go-sdk area:java-sdk area:performance area:ts-sdk kind:task
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.