Split out and handle 'params' in mapped operator #26100

uranusjr · 2022-09-01T08:19:30Z

Allow using params in expanding kwargs. Fix #24014.

Syntax unlocked:

# Two tasks with params={"a": 1} and params={"b": 2}.
MyOperator.partial(task_id="t").expand(params=[{"a": 1}, {"b": 2}])

# Two tasks with params={"a": 0, "b": 1} and params={"a": 0, "b": 2}.
# DAG-level params are also merged. If there are duplicated keys, the
# priority is mapped -> task-partial -> DAG.
MyOperator.partial(task_id="t", params={"a": 0}).expand(params=[{"b": 1}, {"b": 2}])

uranusjr

Some notes explaining the implementation, hopefully this is easier to review with them.

uranusjr · 2022-09-07T22:03:54Z

airflow/models/mappedoperator.py

@@ -162,6 +165,7 @@ class OperatorPartial:
    """

    operator_class: Type["BaseOperator"]
+    params: Union[ParamsDict, dict]


To be able to expand a task against params, we need to be able to “merge” default DAG- and task-level params with the user-mapped params. So params is split out of other partial kwargs and treated specially. Partial kwargs should never contain the "params" key.

uranusjr · 2022-09-07T22:04:42Z

airflow/models/baseoperator.py

@@ -256,7 +256,6 @@ def partial(
    partial_kwargs.setdefault("end_date", end_date)
    partial_kwargs.setdefault("owner", owner)
    partial_kwargs.setdefault("email", email)
-    partial_kwargs.setdefault("params", default_params)


So we don’t set the default params in partial_kwargs…

uranusjr · 2022-09-07T22:04:56Z

airflow/models/baseoperator.py

+    return OperatorPartial(
+        operator_class=operator_class,
+        kwargs=partial_kwargs,
+        params=partial_params,


… but pass them to this separate attribute.

uranusjr · 2022-09-07T22:05:35Z

airflow/decorators/base.py

@@ -363,7 +363,6 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg:
        task_id = get_unique_task_id(partial_kwargs.pop("task_id"), dag, task_group)
        if task_group:
            task_id = task_group.child_id(task_id)
-        params = partial_kwargs.pop("params", None) or default_params


And this line is deleted since it’s now no-op.

uranusjr · 2022-09-07T22:06:39Z

airflow/models/mappedoperator.py

-        # Ordering is significant; mapped kwargs should override partial ones.
+
+        # If params appears in the mapped kwargs, we need to merge it into the
+        # partial params, overriding existing keys.
+        params = copy.copy(self.params)
+        with contextlib.suppress(KeyError):
+            params.update(mapped_kwargs["params"])
+
+        # Ordering is significant; mapped kwargs should override partial ones,
+        # and the specially handled params should be respected.
        return {
            "task_id": self.task_id,
            "dag": self.dag,
            "task_group": self.task_group,
-            "params": self.params,
            "start_date": self.start_date,
            "end_date": self.end_date,
            **self.partial_kwargs,
            **mapped_kwargs,
+            "params": params,
        }


Actual merging happens here. The merged params is put last so it overrides the incomplete mapped_kwargs["params"].

uranusjr · 2022-09-07T22:08:25Z

airflow/models/baseoperator.py

-        """Template all attributes listed in template_fields.
+        """Template all attributes listed in *self.template_fields*.

        This mutates the attributes in-place and is irreversible.

-        :param context: Dict with values to apply on content
-        :param jinja_env: Jinja environment
+        :param context: Context dict with values to apply on content.
+        :param jinja_env: Jinja environment to use for rendering.
        """
        if not jinja_env:
            jinja_env = self.get_template_env()
        self._do_render_template_fields(self, self.template_fields, context, jinja_env, set())
-        return self
+        return None


This does not change any behaviour, I only did it so the behaviour is easier to explain…

uranusjr · 2022-09-07T22:09:58Z

airflow/models/taskinstance.py

+        original_task = self.task
        rendered_task = self.task.render_template_fields(context)
-        if rendered_task is None:  # Compatibility -- custom renderer, assume unmapped.
-            return self.task
-        original_task, self.task = self.task, rendered_task
+        if rendered_task is not None:  # Mapped operator, assign unmapped task.
+            self.task = rendered_task
        return original_task


… here. Previously it’s difficult to follow when rendered_task is None and when it’s not; now it’s simple—a mapped operator (MappedOperator subclases) performs unmapping and returns the unmapped task, while a non-mapped operator (BaseOperator subclasses) return None because no unmapping is needed.

uranusjr · 2022-09-27T09:09:27Z

~~This is going to conflict with #26702 and I want to get that one in first for 2.4.1.~~ Resolved, should be ready for 2.5.0.

It turns out we need to update the template context a bit after unmapping. This also fixes a bug that context["task"] pointed to the mapped task but context["ti"].task is the unmapped task. They now both point to the unmapped one.

boring-cyborg bot added the area:serialization label Sep 1, 2022

uranusjr added area:dynamic-task-mapping AIP-42 and removed area:serialization labels Sep 1, 2022

uranusjr added this to the Airflow 2.4.0 milestone Sep 1, 2022

uranusjr force-pushed the aip-42-task-param-expand branch 5 times, most recently from 5824f87 to 7b2115b Compare September 7, 2022 21:59

uranusjr commented Sep 7, 2022

View reviewed changes

uranusjr marked this pull request as ready for review September 7, 2022 22:49

uranusjr requested review from kaxil, ashb and XD-DENG as code owners September 7, 2022 22:49

ashb modified the milestones: Airflow 2.4.0, Airflow 2.5.0 Sep 8, 2022

uranusjr force-pushed the aip-42-task-param-expand branch from 7b2115b to 50c6989 Compare September 14, 2022 09:18

uranusjr force-pushed the aip-42-task-param-expand branch from 50c6989 to d05e7c2 Compare September 22, 2022 04:35

uranusjr mentioned this pull request Sep 26, 2022

Unable to access operator attrs within Jinja context for mapped tasks #24388

Closed

2 tasks

uranusjr marked this pull request as draft September 27, 2022 09:09

uranusjr added 2 commits September 29, 2022 11:34

Split out and handle 'params' in mapped operator

ba229ad

Implement mapping against 'params'

e9b5861

It turns out we need to update the template context a bit after unmapping. This also fixes a bug that context["task"] pointed to the mapped task but context["ti"].task is the unmapped task. They now both point to the unmapped one.

uranusjr force-pushed the aip-42-task-param-expand branch from d05e7c2 to e9b5861 Compare September 29, 2022 03:42

uranusjr marked this pull request as ready for review September 29, 2022 03:42

dstandish approved these changes Oct 6, 2022

View reviewed changes

uranusjr merged commit 0ec9651 into apache:main Oct 6, 2022

uranusjr deleted the aip-42-task-param-expand branch October 6, 2022 04:47

ephraimbuddy added the type:new-feature Changelog: New Features label Oct 18, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Split out and handle 'params' in mapped operator #26100

Split out and handle 'params' in mapped operator #26100

uranusjr commented Sep 1, 2022 •

edited

Loading

uranusjr left a comment

uranusjr Sep 7, 2022

uranusjr Sep 7, 2022

uranusjr Sep 7, 2022

uranusjr Sep 7, 2022

uranusjr Sep 7, 2022

uranusjr Sep 7, 2022

uranusjr Sep 7, 2022

uranusjr commented Sep 27, 2022 •

edited

Loading

Split out and handle 'params' in mapped operator #26100

Split out and handle 'params' in mapped operator #26100

Conversation

uranusjr commented Sep 1, 2022 • edited Loading

uranusjr left a comment

Choose a reason for hiding this comment

uranusjr Sep 7, 2022

Choose a reason for hiding this comment

uranusjr Sep 7, 2022

Choose a reason for hiding this comment

uranusjr Sep 7, 2022

Choose a reason for hiding this comment

uranusjr Sep 7, 2022

Choose a reason for hiding this comment

uranusjr Sep 7, 2022

Choose a reason for hiding this comment

uranusjr Sep 7, 2022

Choose a reason for hiding this comment

uranusjr Sep 7, 2022

Choose a reason for hiding this comment

uranusjr commented Sep 27, 2022 • edited Loading

uranusjr commented Sep 1, 2022 •

edited

Loading

uranusjr commented Sep 27, 2022 •

edited

Loading