Implement mapped value unpacking

airflow/cli/commands/task_command.py

@@ -224,7 +224,6 @@ def _run_task_by_local_task_job(args, ti):
 
 def _run_raw_task(args, ti: TaskInstance) -> None:
     """Runs the main task handling code"""
-    ti.task = ti.task.unmap()
     ti._run_raw_task(
         mark_success=args.mark_success,
         job_id=args.job_id,
@@ -530,7 +529,6 @@ def task_render(args):
     dag = get_dag(args.subdir, args.dag_id)
     task = dag.get_task(task_id=args.task_id)
     ti = _get_ti(task, args.execution_date_or_run_id, args.map_index, create_if_necessary=True)
-    ti.task = ti.task.unmap()
     ti.render_templates()
     for attr in task.__class__.template_fields:
         print(

airflow/decorators/base.py

@@ -15,12 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import collections.abc
 import functools
 import inspect
-import itertools
 import re
 import sys
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Collection,
@@ -39,19 +40,29 @@
 
 import attr
 import typing_extensions
+from sqlalchemy.orm import Session
 
 from airflow.compat.functools import cache, cached_property
 from airflow.exceptions import AirflowException
 from airflow.models.abstractoperator import DEFAULT_RETRIES, DEFAULT_RETRY_DELAY
 from airflow.models.baseoperator import BaseOperator, coerce_resources, coerce_retry_delay, parse_retries
 from airflow.models.dag import DAG, DagContext
-from airflow.models.mappedoperator import MappedOperator
+from airflow.models.mappedoperator import (
+    MappedOperator,
+    ValidationSource,
+    create_mocked_kwargs,
+    get_mappable_types,
+)
 from airflow.models.pool import Pool
 from airflow.models.xcom_arg import XComArg
 from airflow.typing_compat import Protocol
 from airflow.utils import timezone
 from airflow.utils.context import Context
 from airflow.utils.task_group import TaskGroup, TaskGroupContext
+from airflow.utils.types import NOTSET
+
+if TYPE_CHECKING:
+    from airflow.models.mappedoperator import MapArgument
 
 
 def validate_python_callable(python_callable: Any) -> None:
@@ -269,22 +280,21 @@ def __call__(self, *args, **kwargs) -> XComArg:
             op.doc_md = self.function.__doc__
         return XComArg(op)
 
-    def _validate_arg_names(self, funcname: str, kwargs: Dict[str, Any], valid_names: Set[str] = set()):
-        unknown_args = kwargs.copy()
-        for name in itertools.chain(self.function_arg_names, valid_names):
-            unknown_args.pop(name, None)
-
-            if not unknown_args:
-                # If we have no args left ot check, we are valid
-                return
+    def _validate_arg_names(self, func: ValidationSource, kwargs: Dict[str, Any]):
+        kwargs_left = kwargs.copy()
+        for arg_name in self.function_arg_names:
+            value = kwargs_left.pop(arg_name, NOTSET)
+            if func != "map" or value is NOTSET or isinstance(value, get_mappable_types()):
+                continue
+            raise ValueError(f"{func} got unexpected value{type(value)!r} for keyword argument {arg_name!r}")
 
-        if len(unknown_args) == 1:
-            raise TypeError(f'{funcname} got unexpected keyword argument {next(iter(unknown_args))!r}')
-        else:
-            names = ", ".join(repr(n) for n in unknown_args)
-            raise TypeError(f'{funcname} got unexpected keyword arguments {names}')
+        if len(kwargs_left) == 1:
+            raise TypeError(f"{func} got unexpected keyword argument {next(iter(kwargs_left))!r}")
+        elif kwargs_left:
+            names = ", ".join(repr(n) for n in kwargs_left)
+            raise TypeError(f"{func} got unexpected keyword arguments {names}")
 
-    def map(self, **kwargs) -> XComArg:
+    def map(self, **kwargs: "MapArgument") -> XComArg:
         self._validate_arg_names("map", kwargs)
 
         partial_kwargs = self.kwargs.copy()
@@ -293,7 +303,6 @@ def map(self, **kwargs) -> XComArg:
         task_group = partial_kwargs.pop("task_group", TaskGroupContext.get_current_task_group(dag))
         task_id = get_unique_task_id(partial_kwargs.pop("task_id"), dag, task_group)
         params = partial_kwargs.pop("params", None)
-        partial_op_kwargs = partial_kwargs.pop("op_kwargs", {})
 
         # Logic here should be kept in sync with BaseOperatorMeta.partial().
         if "task_concurrency" in partial_kwargs:
@@ -310,12 +319,14 @@ def map(self, **kwargs) -> XComArg:
         )
         partial_kwargs["resources"] = coerce_resources(partial_kwargs.get("resources"))
         partial_kwargs.setdefault("executor_config", {})
+        partial_kwargs.setdefault("op_args", [])
+        partial_kwargs.setdefault("op_kwargs", {})
 
         # Mypy does not work well with a subclassed attrs class :(
         _MappedOperator = cast(Any, DecoratedMappedOperator)
         operator = _MappedOperator(
             operator_class=self.operator_class,
-            mapped_kwargs={"op_args": [], "op_kwargs": kwargs},
+            mapped_kwargs={},
             partial_kwargs=partial_kwargs,
             task_id=task_id,
             params=params,
@@ -334,7 +345,7 @@ def map(self, **kwargs) -> XComArg:
             end_date=end_date,
             multiple_outputs=self.multiple_outputs,
             python_callable=self.function,
-            partial_op_kwargs=partial_op_kwargs,
+            mapped_op_kwargs=kwargs,
         )
         return XComArg(operator=operator)
 
@@ -347,12 +358,7 @@ def partial(self, **kwargs) -> "_TaskDecorator[Function, OperatorSubclass]":
         return attr.evolve(self, kwargs={**self.kwargs, "op_kwargs": op_kwargs})
 
 
-def _merge_kwargs(
-    kwargs1: Dict[str, XComArg],
-    kwargs2: Dict[str, XComArg],
-    *,
-    fail_reason: str,
-) -> Dict[str, XComArg]:
+def _merge_kwargs(kwargs1: Dict[str, Any], kwargs2: Dict[str, Any], *, fail_reason: str) -> Dict[str, Any]:
     duplicated_keys = set(kwargs1).intersection(kwargs2)
     if len(duplicated_keys) == 1:
         raise TypeError(f"{fail_reason} argument: {duplicated_keys.pop()}")
@@ -369,31 +375,43 @@ class DecoratedMappedOperator(MappedOperator):
     multiple_outputs: bool
     python_callable: Callable
 
-    # We can't save these in partial_kwargs because op_args and op_kwargs need
-    # to be present in mapped_kwargs, and MappedOperator prevents duplication.
-    partial_op_kwargs: Dict[str, Any]
+    # We can't save these in mapped_kwargs because op_kwargs need to be present
+    # in partial_kwargs, and MappedOperator prevents duplication.
+    mapped_op_kwargs: Dict[str, "MapArgument"]
 
     @classmethod
     @cache
     def get_serialized_fields(cls):
-        # The magic argument-less super() does not work well with @cache
-        # (actually lru_cache in general), so we use the explicit form instead.
+        # The magic super() doesn't work here, so we use the explicit form.
+        # Not using super(..., cls) to work around pyupgrade bug.
         sup = super(DecoratedMappedOperator, DecoratedMappedOperator)
-        return sup.get_serialized_fields() | {"partial_op_kwargs"}
+        return sup.get_serialized_fields() | {"mapped_op_kwargs"}
 
-    def _create_unmapped_operator(
-        self,
-        *,
-        mapped_kwargs: Dict[str, Any],
-        partial_kwargs: Dict[str, Any],
-        real: bool,
-    ) -> "BaseOperator":
+    def __attrs_post_init__(self):
+        # The magic super() doesn't work here, so we use the explicit form.
+        # Not using super(..., self) to work around pyupgrade bug.
+        super(DecoratedMappedOperator, DecoratedMappedOperator).__attrs_post_init__(self)
+        XComArg.apply_upstream_relationship(self, self.mapped_op_kwargs)
+
+    def _get_expansion_kwargs(self) -> Dict[str, "MapArgument"]:
+        """The kwargs to calculate expansion length against.
+
+        Different from classic operators, a decorated (taskflow) operator's
+        ``map()`` contributes to the ``op_kwargs`` operator argument (not the
+        operator arguments themselves), and should therefore expand against it.
+        """
+        return self.mapped_op_kwargs
+
+    def _create_unmapped_operator(self, *, mapped_kwargs: Dict[str, Any], real: bool) -> "BaseOperator":
         assert not isinstance(self.operator_class, str)
-        mapped_kwargs = mapped_kwargs.copy()
-        del mapped_kwargs["op_kwargs"]
+        partial_kwargs = self.partial_kwargs.copy()
+        if real:
+            mapped_op_kwargs: Dict[str, Any] = self.mapped_op_kwargs
+        else:
+            mapped_op_kwargs = create_mocked_kwargs(self.mapped_op_kwargs)
         op_kwargs = _merge_kwargs(
-            self.partial_op_kwargs,
-            self.mapped_kwargs["op_kwargs"],  # We want to "original" op_kwargs.
+            partial_kwargs.pop("op_kwargs"),
+            mapped_op_kwargs,
             fail_reason="mapping already partial",
         )
         return self.operator_class(
@@ -408,6 +426,14 @@ def _create_unmapped_operator(
             **mapped_kwargs,
         )
 
+    def _expand_mapped_field(self, key: str, content: Any, context: Context, *, session: Session) -> Any:
+        if key != "op_kwargs" or not isinstance(content, collections.abc.Mapping):
+            return content
+        # The magic super() doesn't work here, so we use the explicit form.
+        # Not using super(..., self) to work around pyupgrade bug.
+        sup: Any = super(DecoratedMappedOperator, DecoratedMappedOperator)
+        return {k: sup._expand_mapped_field(self, k, v, context, session=session) for k, v in content.items()}
+
 
 class Task(Generic[Function]):
     """Declaration of a @task-decorated callable for type-checking.
@@ -424,7 +450,7 @@ class Task(Generic[Function]):
 
     function: Function
 
-    def map(self, **kwargs: Any) -> XComArg:
+    def map(self, **kwargs: "MapArgument") -> XComArg:
         ...
 
     def partial(self, **kwargs: Any) -> "Task[Function]":

airflow/executors/debug_executor.py

@@ -76,7 +76,6 @@ def _run_task(self, ti: TaskInstance) -> bool:
         key = ti.key
         try:
             params = self.tasks_params.pop(ti.key, {})
-            ti.task = ti.task.unmap()
             ti._run_raw_task(job_id=ti.job_id, **params)
             self.change_state(key, State.SUCCESS)
             ti._run_finished_callback()

airflow/jobs/backfill_job.py

@@ -265,7 +265,7 @@ def _manage_executor_state(
             if ti.state not in self.STATES_COUNT_AS_RUNNING:
                 for node in ti.task.mapped_dependants():
                     assert isinstance(node, MappedOperator)
-                    yield node, ti.run_id, node.expand_mapped_task(ti, session)
+                    yield node, ti.run_id, node.expand_mapped_task(ti.run_id, session=session)
 
     @provide_session
     def _get_dag_run(self, dagrun_info: DagRunInfo, dag: DAG, session: Session = None):