Add TaskMap and TaskInstance.map_id

airflow/decorators/base.py

@@ -277,24 +277,18 @@ def _validate_arg_names(self, funcname: str, kwargs: Dict[str, Any], valid_names
     def map(
         self, *, dag: Optional["DAG"] = None, task_group: Optional["TaskGroup"] = None, **kwargs
     ) -> XComArg:
-
+        self._validate_arg_names("map", kwargs)
         dag = dag or DagContext.get_current_dag()
         task_group = task_group or TaskGroupContext.get_current_task_group(dag)
         task_id = get_unique_task_id(self.kwargs['task_id'], dag, task_group)
 
-        self._validate_arg_names("map", kwargs)
-
-        operator = MappedOperator(
-            operator_class=self.operator_class,
-            task_id=task_id,
+        operator = MappedOperator.from_decorator(
+            decorator=self,
             dag=dag,
             task_group=task_group,
-            partial_kwargs=self.kwargs,
-            # Set them to empty to bypass the validation, as for decorated stuff we validate ourselves
-            mapped_kwargs={},
+            task_id=task_id,
+            mapped_kwargs=kwargs,
         )
-        operator.mapped_kwargs.update(kwargs)
-
         return XComArg(operator=operator)
 
     def partial(

airflow/exceptions.py

@@ -99,6 +99,17 @@ class AirflowFailException(AirflowException):
     """Raise when the task should be failed without retrying."""
 
 
+class UnmappableXComPushed(AirflowException):
+    """Raise when an unmappable value is pushed as a mapped downstream's dependency."""
+
+    def __init__(self, value: Any) -> None:
+        super().__init__(value)
+        self.value = value
+
+    def __str__(self) -> str:
+        return f"unmappable return type {type(self.value).__qualname__!r}"
+
+
 class AirflowDagCycleException(AirflowException):
     """Raise when there is a cycle in DAG definition."""
 

airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py

@@ -0,0 +1,120 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Add TaskMap and map_index on TaskInstance.
+
+Revision ID: e655c0453f75
+Revises: 587bdf053233
+Create Date: 2021-12-13 22:59:41.052584
+"""
+
+from alembic import op
+from sqlalchemy import Column, ForeignKeyConstraint, Integer
+
+from airflow.models.base import StringID
+from airflow.utils.sqlalchemy import ExtendedJSON
+
+# Revision identifiers, used by Alembic.
+revision = "e655c0453f75"
+down_revision = "587bdf053233"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    """Add TaskMap and map_index on TaskInstance."""
+    # We need to first remove constraints on task_reschedule since they depend on task_instance.
+    with op.batch_alter_table("task_reschedule") as batch_op:
+        batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey")
+        batch_op.drop_index("idx_task_reschedule_dag_task_run")
+
+    # Change task_instance's primary key.
+    with op.batch_alter_table("task_instance") as batch_op:
+        # I think we always use this name for TaskInstance after 7b2661a43ba3?
+        batch_op.drop_constraint("task_instance_pkey", type_="primary")
+        batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1))
+        batch_op.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id", "map_index"])
+
+    # Re-create task_reschedule's constraints.
+    with op.batch_alter_table("task_reschedule") as batch_op:
+        batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1))
+        batch_op.create_foreign_key(
+            "task_reschedule_ti_fkey",
+            "task_instance",
+            ["dag_id", "task_id", "run_id", "map_index"],
+            ["dag_id", "task_id", "run_id", "map_index"],
+            ondelete="CASCADE",
+        )
+        batch_op.create_index(
+            "idx_task_reschedule_dag_task_run",
+            ["dag_id", "task_id", "run_id", "map_index"],
+            unique=False,
+        )
+
+    # Create task_map.
+    op.create_table(
+        "task_map",
+        Column("dag_id", StringID(), primary_key=True),
+        Column("task_id", StringID(), primary_key=True),
+        Column("run_id", StringID(), primary_key=True),
+        Column("map_index", Integer, primary_key=True),
+        Column("length", Integer, nullable=False),
+        Column("keys", ExtendedJSON, nullable=True),
+        ForeignKeyConstraint(
+            ["dag_id", "task_id", "run_id", "map_index"],
+            [
+                "task_instance.dag_id",
+                "task_instance.task_id",
+                "task_instance.run_id",
+                "task_instance.map_index",
+            ],
+            name="task_map_task_instance_fkey",
+            ondelete="CASCADE",
+        ),
+    )
+
+
+def downgrade():
+    """Remove TaskMap and map_index on TaskInstance."""
+    op.drop_table("task_map")
+
+    with op.batch_alter_table("task_reschedule") as batch_op:
+        batch_op.drop_constraint("task_reschedule_ti_fkey", "foreignkey")
+        batch_op.drop_index("idx_task_reschedule_dag_task_run")
+        batch_op.drop_column("map_index")
+
+    op.execute("DELETE FROM task_instance WHERE map_index != -1")
+
+    with op.batch_alter_table("task_instance") as batch_op:
+        batch_op.drop_constraint("task_instance_pkey", type_="primary")
+        batch_op.drop_column("map_index")
+        batch_op.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id"])
+
+    with op.batch_alter_table("task_reschedule") as batch_op:
+        batch_op.create_foreign_key(
+            "task_reschedule_ti_fkey",
+            "task_instance",
+            ["dag_id", "task_id", "run_id"],
+            ["dag_id", "task_id", "run_id"],
+            ondelete="CASCADE",
+        )
+        batch_op.create_index(
+            "idx_task_reschedule_dag_task_run",
+            ["dag_id", "task_id", "run_id"],
+            unique=False,
+        )

airflow/models/baseoperator.py

@@ -78,6 +78,7 @@
 from airflow.utils.weight_rule import WeightRule
 
 if TYPE_CHECKING:
+    from airflow.decorators.base import _TaskDecorator
     from airflow.models.dag import DAG
     from airflow.utils.task_group import TaskGroup
 
@@ -1632,6 +1633,39 @@ def defer(
     def map(self, **kwargs) -> "MappedOperator":
         return MappedOperator.from_operator(self, kwargs)
 
+    def has_mapped_dependants(self) -> bool:
+        """Whether any downstream dependencies depend on this task for mapping.
+
+        For now, this walks the entire DAG to find mapped nodes that has this
+        current task as an upstream. We cannot use ``downstream_list`` since it
+        only contains operators, not task groups. In the future, we should
+        provide a way to record an DAG node's all downstream nodes instead.
+        """
+        from airflow.utils.task_group import MappedTaskGroup, TaskGroup
+
+        if not self.has_dag():
+            return False
+
+        def _walk_group(group: TaskGroup) -> Iterable[Tuple[str, DAGNode]]:
+            """Recursively walk children in a task group.
+
+            This yields all direct children (including both tasks and task
+            groups), and all children of any task groups.
+            """
+            for key, child in group.children.items():
+                yield key, child
+                if isinstance(child, TaskGroup):
+                    yield from _walk_group(child)
+
+        for key, child in _walk_group(self.dag.task_group):
+            if key == self.task_id:
+                continue
+            if not isinstance(child, (MappedOperator, MappedTaskGroup)):
+                continue
+            if self.task_id in child.upstream_task_ids:
+                return True
+        return False
+
 
 def _validate_kwarg_names_for_mapping(cls: Type[BaseOperator], func_name: str, value: Dict[str, Any]):
     if isinstance(str, cls):
@@ -1668,6 +1702,7 @@ class MappedOperator(DAGNode):
     """Object representing a mapped operator in a DAG"""
 
     operator_class: Type[BaseOperator] = attr.ib(repr=lambda c: c.__name__)
+    task_type: str = attr.ib()
     task_id: str
     partial_kwargs: Dict[str, Any]
     mapped_kwargs: Dict[str, Any] = attr.ib(
@@ -1678,9 +1713,12 @@ class MappedOperator(DAGNode):
     downstream_task_ids: Set[str] = attr.ib(factory=set, repr=False)
 
     task_group: Optional["TaskGroup"] = attr.ib(repr=False)
+
     # BaseOperator-like interface -- needed so we can add oursleves to the dag.tasks
     start_date: Optional[pendulum.DateTime] = attr.ib(repr=False, default=None)
     end_date: Optional[pendulum.DateTime] = attr.ib(repr=False, default=None)
+    owner: str = attr.ib(repr=False, default=conf.get("operators", "DEFAULT_OWNER"))
+    max_active_tis_per_dag: Optional[int] = attr.ib(default=None)
 
     @classmethod
     def from_operator(cls, operator: BaseOperator, mapped_kwargs: Dict[str, Any]) -> "MappedOperator":
@@ -1699,15 +1737,47 @@ def from_operator(cls, operator: BaseOperator, mapped_kwargs: Dict[str, Any]) ->
             end_date=operator.end_date,
             partial_kwargs=operator._BaseOperator__init_kwargs,  # type: ignore
             mapped_kwargs=mapped_kwargs,
+            owner=operator.owner,
+            max_active_tis_per_dag=operator.max_active_tis_per_dag,
         )
 
+    @classmethod
+    def from_decorator(
+        cls,
+        *,
+        decorator: "_TaskDecorator",
+        dag: Optional["DAG"],
+        task_group: Optional["TaskGroup"],
+        task_id: str,
+        mapped_kwargs: Dict[str, Any],
+    ) -> "MappedOperator":
+        """Create a mapped operator from a task decorator.
+
+        Different from ``from_operator``, this DOES NOT validate ``mapped_kwargs``.
+        The task decorator calling this should be responsible for validation.
+        """
+        operator = MappedOperator(
+            operator_class=decorator.operator_class,
+            partial_kwargs=decorator.kwargs,
+            mapped_kwargs={},
+            task_id=task_id,
+            dag=dag,
+            task_group=task_group,
+        )
+        operator.mapped_kwargs.update(mapped_kwargs)
+        return operator
+
     def __attrs_post_init__(self):
         if self.task_group:
             self.task_id = self.task_group.child_id(self.task_id)
             self.task_group.add(self)
         if self.dag:
             self.dag.add_task(self)
 
+    @task_type.default
+    def _default_task_type(self):
+        return self.operator_class.__name__
+
     @task_group.default
     def _default_task_group(self):
         from airflow.utils.task_group import TaskGroupContext

airflow/models/taskinstance.py

@@ -15,6 +15,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import collections.abc
 import contextlib
 import hashlib
 import logging
@@ -82,11 +83,13 @@
     DagRunNotFound,
     TaskDeferralError,
     TaskDeferred,
+    UnmappableXComPushed,
 )
 from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
 from airflow.models.log import Log
 from airflow.models.param import ParamsDict
 from airflow.models.taskfail import TaskFail
+from airflow.models.taskmap import TaskMap
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.xcom import XCOM_RETURN_KEY, XCom
 from airflow.plugins_manager import integrate_macros_plugins
@@ -340,6 +343,8 @@ class TaskInstance(Base, LoggingMixin):
     task_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False)
     dag_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False)
     run_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False)
+    map_index = Column(Integer, primary_key=True, nullable=False, default=-1)
+
     start_date = Column(UtcDateTime)
     end_date = Column(UtcDateTime)
     duration = Column(Float)
@@ -424,10 +429,12 @@ def __init__(
         execution_date: Optional[datetime] = None,
         run_id: Optional[str] = None,
         state: Optional[str] = None,
+        map_index: int = -1,
     ):
         super().__init__()
         self.dag_id = task.dag_id
         self.task_id = task.task_id
+        self.map_index = map_index
         self.refresh_from_task(task)
         self._log = logging.getLogger("airflow.task")
 
@@ -760,6 +767,7 @@ def refresh_from_db(self, session=NEW_SESSION, lock_for_update=False) -> None:
             TaskInstance.dag_id == self.dag_id,
             TaskInstance.task_id == self.task_id,
             TaskInstance.run_id == self.run_id,
+            TaskInstance.map_index == self.map_index,
         )
 
         if lock_for_update:
@@ -1537,7 +1545,9 @@ def _execute_task(self, context, task_copy):
             result = execute_callable(context=context)
         # If the task returns a result, push an XCom containing it
         if task_copy.do_xcom_push and result is not None:
-            self.xcom_push(key=XCOM_RETURN_KEY, value=result)
+            with create_session() as session:
+                self.xcom_push(key=XCOM_RETURN_KEY, value=result, session=session)
+                self._record_task_map_for_downstreams(result, session=session)
         return result
 
     @provide_session
@@ -2128,6 +2138,14 @@ def set_duration(self) -> None:
             self.duration = None
         self.log.debug("Task Duration set to %s", self.duration)
 
+    def _record_task_map_for_downstreams(self, value: Any, *, session: Session) -> None:
+        if not self.task.has_mapped_dependants():
+            return
+        if not isinstance(value, collections.abc.Collection) or isinstance(value, (bytes, str)):
+            self.log.info("Failing %s for unmappable XCom push %r", self.key, type(value).__qualname__)
+            raise UnmappableXComPushed(value)
+        session.merge(TaskMap.from_task_instance_xcom(self, value))
+
     @provide_session
     def xcom_push(
         self,