dstandish commented on code in PR #69565:
URL: https://github.com/apache/airflow/pull/69565#discussion_r3580418928
##########
airflow-core/tests/unit/models/test_taskinstance.py:
##########
@@ -3291,6 +3291,197 @@ def show(value):
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) and are attached to the session.
+ """
+ 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")
+ 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.
Review Comment:
Oh i misunderstood what you were suggesting
i thought you were saying do
set_committed_value(ti, 'map_index', 0)
so that's why i was liike "that doesn't seem like it would save a select"
but you were saying `set_committed_value(ti, "dag_run", dr)`
it turns out that on line `taskmap.py:203`, `dr = unmapped_ti.dag_run`
accesses the relationship on the unmapped TI, which lazy-loads and caches it on
the instance. That access always happens before the TI is appended to the
results (`taskmap.py:216`), so by the time dependency evaluation calls
`ti.get_dagrun()`, it's a cache hit — `get_dagrun` returns the cached
`self.dag_run` without querying.
We still *could* add set_committed_value there -- it wouldn't hurt -- but
because the DR is already there, it is not strictly necessary. LMKWYT
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
