[GitHub] [airflow] uranusjr commented on a diff in pull request #27876: Make MappedTaskGroup depend on its expand inputs

[GitBox <gi...@apache.org>]

uranusjr commented on code in PR #27876:
URL: https://github.com/apache/airflow/pull/27876#discussion_r1031050752


##########
airflow/utils/task_group.py:
##########
@@ -507,6 +508,15 @@ class MappedTaskGroup(TaskGroup):
     def __init__(self, *, expand_input: ExpandInput, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self._expand_input = expand_input
+        for op, _ in expand_input.iter_references():
+            self.set_upstream(op)
+
+    def iter_mapped_dependencies(self) -> Iterator[Operator]:
+        """Upstream dependencies that provide XComs used by this mapped task group."""
+        from airflow.models.xcom_arg import XComArg
+
+        for op, _ in XComArg.iter_xcom_references(self._expand_input):
+            yield op

Review Comment:
   Implemented to match MappedOperator, so `iter_mapped_dependants` can stay unchanged (except the type hints).



##########
airflow/utils/task_group.py:
##########
@@ -507,6 +508,15 @@ class MappedTaskGroup(TaskGroup):
     def __init__(self, *, expand_input: ExpandInput, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self._expand_input = expand_input
+        for op, _ in expand_input.iter_references():
+            self.set_upstream(op)

Review Comment:
   This is the key fix.



##########
airflow/jobs/backfill_job.py:
##########
@@ -271,10 +271,20 @@ def _manage_executor_state(
                 self.log.error(msg)
                 ti.handle_failure(error=msg)
                 continue
+
+            def _iter_task_needing_expansion() -> Iterator[AbstractOperator]:
+                from airflow.models.mappedoperator import AbstractOperator
+
+                for node in self.dag.get_task(ti.task_id, include_subdags=True).iter_mapped_dependants():
+                    if isinstance(node, AbstractOperator):
+                        yield node
+                    else:  # A (mapped) task group. All its children need expansion.
+                        yield from node.iter_tasks()
+
             if ti.state not in self.STATES_COUNT_AS_RUNNING:
                 # Don't use ti.task; if this task is mapped, that attribute
                 # would hold the unmapped task. We need to original task here.
-                for node in self.dag.get_task(ti.task_id, include_subdags=True).iter_mapped_dependants():
+                for node in _iter_task_needing_expansion():

Review Comment:
   Bug caught by cascaded type hint updates.



##########
airflow/models/taskinstance.py:
##########
@@ -2226,18 +2239,14 @@ def _record_task_map_for_downstreams(self, task: Operator, value: Any, *, sessio
             return
         # TODO: We don't push TaskMap for mapped task instances because it's not
         # currently possible for a downstream to depend on one individual mapped
-        # task instance. This will change when we implement task group mapping,
-        # and we'll need to further analyze the mapped task case.
+        # task instance. This will change when we implement task mapping inside
+        # a mapped task group, and we'll need to further analyze the case.

Review Comment:
   This comment is fixed to reflect the up-to-date exact reason.



##########
airflow/models/taskinstance.py:
##########
@@ -2226,18 +2239,14 @@ def _record_task_map_for_downstreams(self, task: Operator, value: Any, *, sessio
             return
         # TODO: We don't push TaskMap for mapped task instances because it's not
         # currently possible for a downstream to depend on one individual mapped
-        # task instance. This will change when we implement task group mapping,
-        # and we'll need to further analyze the mapped task case.
+        # task instance. This will change when we implement task mapping inside
+        # a mapped task group, and we'll need to further analyze the case.
         if task.is_mapped:
             return
         if value is None:
             raise XComForMappingNotPushed()
-        if not isinstance(value, (collections.abc.Sequence, dict)):
-            raise UnmappableXComTypePushed(value)
-        if isinstance(value, (bytes, str)):
+        if not _is_mappable_value(value):
             raise UnmappableXComTypePushed(value)
-        if TYPE_CHECKING:  # The isinstance() checks above guard this.
-            assert isinstance(value, collections.abc.Collection)

Review Comment:
   Just some refactorings to use PEP 647 TypeGuard.



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org