dstandish commented on code in PR #69565:
URL: https://github.com/apache/airflow/pull/69565#discussion_r3545759514
##########
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:
i am *not sure* if you would want to. when you set dag_run you are setting
a relationship attr and, it's easy to imagine how that could result in extra
selects. but with just an attr, i suspect it would not do anything beyond
emitting the DML when session flushed, which, if you're changing the index, it
would have to do anyway. lmk
--
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]
