This is an automated email from the ASF dual-hosted git repository.
dstandish pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new cad6220034e Speed up dynamic task mapping expansion (#69565)
cad6220034e is described below
commit cad6220034e87c586eb87e14b5a75f4b82984fd0
Author: Daniel Standish <[email protected]>
AuthorDate: Tue Jul 14 10:09:24 2026 -0700
Speed up dynamic task mapping expansion (#69565)
Expanding a mapped task creates the N downstream task instances inside
`DagRun.update_state`. Two `N+1` query patterns there turned a wide fan-out
into a scheduler stall that scaled linearly with map width: a per-index
`session.merge()` to persist each new instance, and a per-index `dag_run`
`SELECT` when dependency evaluation later called `get_dagrun()` on instances
whose dag_run relationship was never loaded.
Persist the instances with a batched `session.add()`/`flush()` and prime the
shared `dag_run` relationship once during expansion, so both scale with a
constant number of queries instead of with map width.
We could have also defensively added a `set_committed_value` on the dagrun
attr for the index-0 case but it's not necessary as is because the dag run
object is cached for that reused TI object.
---
airflow-core/src/airflow/models/dagrun.py | 18 +--
airflow-core/src/airflow/models/taskinstance.py | 21 +++
airflow-core/src/airflow/models/taskmap.py | 16 +-
.../tests/unit/models/test_taskinstance.py | 164 +++++++++++++++++++++
4 files changed, 203 insertions(+), 16 deletions(-)
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index e45a990a3fb..a51f16e51f6 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -79,7 +79,7 @@ from airflow.models import Deadline, Log
from airflow.models.backfill import Backfill
from airflow.models.base import Base, StringID
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
-from airflow.models.taskinstance import TaskInstance as TI,
clear_task_instances
+from airflow.models.taskinstance import TaskInstance as TI,
_add_and_prime_mapped_ti, clear_task_instances
from airflow.models.taskinstancehistory import TaskInstanceHistory as TIH
from airflow.models.tasklog import LogTemplate
from airflow.models.taskmap import TaskMap
@@ -2045,7 +2045,7 @@ class DagRun(Base, LoggingMixin):
def _revise_map_indexes_if_mapped(
self, task: Operator, *, dag_version_id: UUID | None, session: Session
- ) -> Iterator[TI]:
+ ) -> list[TI]:
"""
Check if task increased or reduced in length and handle appropriately.
@@ -2056,14 +2056,13 @@ class DagRun(Base, LoggingMixin):
"""
from airflow.models.expandinput import NotFullyPopulated
from airflow.serialization.definitions.mappedoperator import
get_mapped_ti_count
- from airflow.settings import task_instance_mutation_hook
try:
total_length = get_mapped_ti_count(task, self.run_id,
session=session)
except NotMapped:
- return # Not a mapped task, don't need to do anything.
+ return [] # Not a mapped task, don't need to do anything.
except NotFullyPopulated:
- return # Upstreams not ready, don't need to revise this yet.
+ return [] # Upstreams not ready, don't need to revise this yet.
query = session.scalars(
select(TI.map_index).where(
@@ -2088,16 +2087,17 @@ class DagRun(Base, LoggingMixin):
)
session.flush()
+ new_tis: list[TI] = []
for index in range(total_length):
if index in existing_indexes:
continue
ti = TI(task, run_id=self.run_id, map_index=index, state=None,
dag_version_id=dag_version_id)
self.log.debug("Expanding TIs upserted %s", ti)
- task_instance_mutation_hook(ti, dag_run=self)
- ti = session.merge(ti)
- ti.refresh_from_task(task, dag_run=self)
+ _add_and_prime_mapped_ti(ti, task, self, session=session)
+ new_tis.append(ti)
+ if new_tis:
session.flush()
- yield ti
+ return new_tis
@classmethod
@provide_session
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index 9f8eb0805c4..1c7de2185e6 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -202,6 +202,27 @@ def _stop_remaining_tasks(*, task_instance: TaskInstance,
task_teardown_map=None
log.info("Not skipping teardown task '%s'", ti.task_id)
+def _add_and_prime_mapped_ti(
+ ti: TaskInstance,
+ task: Operator,
+ dag_run: DagRun,
+ *,
+ session: Session,
+ context_carrier: dict | None = None,
+) -> None:
+ """
+ Attach a newly-created mapped TI to the session and prime its ``dag_run``
cache.
+
+ :meta private:
+ """
+ task_instance_mutation_hook(ti, dag_run=dag_run)
+ session.add(ti)
+ if context_carrier is not None:
+ ti.context_carrier = context_carrier
+ ti.refresh_from_task(task, dag_run=dag_run)
+ set_committed_value(ti, "dag_run", dag_run)
+
+
def _recalculate_dagrun_queued_at_deadlines(
dagrun: DagRun, new_queued_at: datetime, session: Session
) -> None:
diff --git a/airflow-core/src/airflow/models/taskmap.py
b/airflow-core/src/airflow/models/taskmap.py
index b7c394a7c13..96b3c0831cb 100644
--- a/airflow-core/src/airflow/models/taskmap.py
+++ b/airflow-core/src/airflow/models/taskmap.py
@@ -141,7 +141,7 @@ class TaskMap(TaskInstanceDependencies):
order by map index, and the maximum map index value.
"""
from airflow.models.expandinput import NotFullyPopulated
- from airflow.models.taskinstance import TaskInstance
+ from airflow.models.taskinstance import TaskInstance,
_add_and_prime_mapped_ti
from airflow.serialization.definitions.baseoperator import
SerializedBaseOperator
from airflow.serialization.definitions.mappedoperator import (
SerializedMappedOperator,
@@ -256,8 +256,8 @@ class TaskMap(TaskInstanceDependencies):
)
)
+ new_tis: list[TaskInstance] = []
for index in indexes_to_map:
- # TODO: Make more efficient with
bulk_insert_mappings/bulk_save_mappings.
ti = TaskInstance(
task,
run_id=run_id,
@@ -266,11 +266,13 @@ class TaskMap(TaskInstanceDependencies):
dag_version_id=dag_version_id,
)
task.log.debug("Expanding TIs upserted %s", ti)
- task_instance_mutation_hook(ti, dag_run=dr)
- ti = session.merge(ti)
- ti.context_carrier = new_task_run_carrier(dr.context_carrier)
- ti.refresh_from_task(task, dag_run=dr) # session.merge() loses
task information.
- all_expanded_tis.append(ti)
+ _add_and_prime_mapped_ti(
+ ti, task, dr, session=session,
context_carrier=new_task_run_carrier(dr.context_carrier)
+ )
+ new_tis.append(ti)
+ if new_tis:
+ session.flush()
+ all_expanded_tis.extend(new_tis)
# Coerce the None case to 0 -- these two are almost treated
identically,
# except the unmapped ti (if exists) is marked to different states.
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index f2f5675ae56..7f1a193d963 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -3289,6 +3289,170 @@ class TestMappedTaskInstanceReceiveValue:
dag_maker.run_ti(ti.task_id, dag_run=dag_run,
map_index=ti.map_index, session=session)
assert outputs == expected_outputs
+ def test_map_xcom_wide_batched_expand(self, dag_maker, session):
+ """Wide XCom-driven expand goes through the batched add_all()/flush()
path.
+
+ Exercises ``TaskMap.expand_mapped_task`` over a 20-element upstream
XCom and
+ asserts the batched expansion creates exactly N mapped TIs with
contiguous
+ ``map_index`` 0..N-1, the expected ``None`` (schedulable) state, and
that the
+ returned instances are usable: they keep their ``.task`` (no merge()
that drops
+ it), are attached to the session, and have ``dag_run`` primed so a
later
+ ``ti.get_dagrun()`` -- as dependency evaluation makes per index -- is
a cache hit
+ rather than an N+1 SELECT. Also pins the expansion call's query count.
+ """
+ from sqlalchemy import event
+ from sqlalchemy.orm.base import NO_VALUE
+
+ width = 20
+ upstream_return = list(range(width))
+
+ with dag_maker(dag_id="xcom_wide", session=session, serialized=True)
as dag:
+
+ @dag.task
+ def emit():
+ return upstream_return
+
+ @dag.task
+ def show(value):
+ return value
+
+ show.expand(value=emit())
+
+ dag_run = dag_maker.create_dagrun()
+ emit_ti = dag_run.get_task_instance("emit", session=session)
+ emit_ti.refresh_from_task(dag_maker.serialized_dag.get_task("emit"))
+ dag_maker.run_ti(emit_ti.task_id, dag_run=dag_run, session=session)
+
+ show_task = dag_maker.serialized_dag.get_task("show")
+ # Pins the query count so a regression back to per-index
session.merge() -- which
+ # would issue a merge-load + reload SELECT per index -- fails this
test, not just
+ # a slower one. Measured at 7 for this fixture; margin allows for
minor backend
+ # differences while staying far below what a per-index merge() would
cost.
+ with assert_queries_count(7, margin=2):
+ mapped_tis, max_map_index = TaskMap.expand_mapped_task(show_task,
dag_run.run_id, session=session)
+
+ # Correct count + contiguous indexes 0..N-1.
+ assert len(mapped_tis) == width
+ assert max_map_index + 1 == width
+ assert sorted(ti.map_index for ti in mapped_tis) == list(range(width))
+
+ # Freshly-expanded mapped TIs are schedulable (state None) and are
attached to
+ # the session. The batched path (indexes >= 1) additionally keeps its
``.task``
+ # because we never merge(); index 0 is the repurposed unmapped TI
(loaded from
+ # the DB, so ``.task`` is None) which is stock behaviour untouched by
this change.
+ for ti in mapped_tis:
+ assert ti.state is None
+ assert ti in session
+ batched_tis = [ti for ti in mapped_tis if ti.map_index >= 1]
+ assert len(batched_tis) == width - 1
+ for ti in batched_tis:
+ assert ti.task is not None
+ assert sa_inspect(ti).attrs.dag_run.loaded_value is not NO_VALUE
+
+ # And they are actually persisted as rows.
+ persisted = session.scalars(
+ select(TI)
+ .where(TI.task_id == "show", TI.dag_id == dag_run.dag_id,
TI.run_id == dag_run.run_id)
+ .order_by(TI.map_index)
+ ).all()
+ assert [ti.map_index for ti in persisted] == list(range(width))
+
+ # dag_run priming means get_dagrun() is a cache hit: zero SELECTs
against dag_run.
+ dag_run_selects = []
+
+ def _on_cursor_execute(conn, cursor, statement, parameters, context,
executemany):
+ normalized = " ".join(statement.lower().split())
+ if "select" in normalized and "from dag_run" in normalized:
+ dag_run_selects.append(statement)
+
+ event.listen(session.bind, "after_cursor_execute", _on_cursor_execute)
+ try:
+ for ti in batched_tis:
+ returned = ti.get_dagrun(session=session)
+ assert returned.run_id == dag_run.run_id
+ finally:
+ event.remove(session.bind, "after_cursor_execute",
_on_cursor_execute)
+ assert dag_run_selects == [], (
+ f"get_dagrun() on primed mapped TIs issued {len(dag_run_selects)}
dag_run SELECT(s); "
+ "the relationship priming should make these cache hits"
+ )
+
+ def test_revise_map_indexes_grow_batched(self, dag_maker, session):
+ """``DagRun._revise_map_indexes_if_mapped`` adds the new TIs via the
batched path.
+
+ Simulates a mid-run length grow: expand once at a narrow width, then
re-run the
+ revise path after the upstream XCom has grown, and assert the
additional indexes
+ are created contiguously without losing any existing ones. Also pins
the revise
+ call's query count.
+ """
+ from airflow.models.xcom import XComModel
+
+ with dag_maker(dag_id="xcom_revise", session=session, serialized=True)
as dag:
+
+ @dag.task
+ def emit():
+ return [1, 2, 3]
+
+ @dag.task
+ def show(value):
+ return value
+
+ show.expand(value=emit())
+
+ dag_run = dag_maker.create_dagrun()
+ emit_ti = dag_run.get_task_instance("emit", session=session)
+ emit_ti.refresh_from_task(dag_maker.serialized_dag.get_task("emit"))
+ dag_maker.run_ti(emit_ti.task_id, dag_run=dag_run, session=session)
+
+ show_task = dag_maker.serialized_dag.get_task("show")
+ mapped_tis, max_map_index = TaskMap.expand_mapped_task(show_task,
dag_run.run_id, session=session)
+ assert len(mapped_tis) == 3
+ assert max_map_index == 2
+
+ # Grow the upstream's pushed length 3 -> 5 by rewriting the
return-value XCom and
+ # the TaskMap row that records the mapped length.
+ XComModel.set(
+ key="return_value",
+ value=[1, 2, 3, 4, 5],
+ dag_id=dag_run.dag_id,
+ task_id="emit",
+ run_id=dag_run.run_id,
+ session=session,
+ )
+ task_map = session.scalars(
+ select(TaskMap).where(
+ TaskMap.dag_id == dag_run.dag_id,
+ TaskMap.task_id == "emit",
+ TaskMap.run_id == dag_run.run_id,
+ )
+ ).one()
+ task_map.length = 5
+ task_map.keys = None
+ session.flush()
+
+ # Pins the query count so a regression back to per-index
session.merge() -- which
+ # would issue a merge-load + reload SELECT per index -- fails this
test. Measured
+ # at 3 for this fixture (2 new indexes); margin allows for minor
backend
+ # differences while staying below what a per-index merge() would cost.
+ with assert_queries_count(3, margin=1):
+ new_tis = list(
+ dag_run._revise_map_indexes_if_mapped(
+ show_task, dag_version_id=mapped_tis[0].dag_version_id,
session=session
+ )
+ )
+ # Only the two brand-new indexes (3, 4) are created, contiguous and
usable.
+ assert sorted(ti.map_index for ti in new_tis) == [3, 4]
+ for ti in new_tis:
+ assert ti.task is not None
+ assert ti in session
+
+ persisted = session.scalars(
+ select(TI)
+ .where(TI.task_id == "show", TI.dag_id == dag_run.dag_id,
TI.run_id == dag_run.run_id)
+ .order_by(TI.map_index)
+ ).all()
+ assert [ti.map_index for ti in persisted] == [0, 1, 2, 3, 4]
+
def test_map_literal_cross_product(self, dag_maker, session):
"""Test a mapped task with literal cross product args expand
properly."""
outputs = []
