[airflow] branch main updated: Documentation on task mapping additions (#24489)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 3221ed5968 Documentation on task mapping additions (#24489)
3221ed5968 is described below

commit 3221ed5968423ea7a0dc7e1a4b51084351c2d56b
Author: Tzu-ping Chung <ur...@gmail.com>
AuthorDate: Sat Aug 6 00:56:47 2022 +0800

    Documentation on task mapping additions (#24489)
    
    * Documentation on task mapping additions
    
    None of these actually exist yet. These are the additions planned for
    2.5 on AIP-42 (task mapping), namely expand_kwargs(), map(), filter
    (using Airflow's existing mechanism), and zip().
    
    The examples for expand_kwargs() and zip() are pretty bad now.
    Suggestions are much welcomed.
    
    Let's discuss the interface in this patch's thread; once we settle on
    the syntax and semantic, I'll start implementing them.
    
    * Detect common user error passing a task to map()
    
    The argument to map() must be a plain function, not a task. This hook
    checks for that.
---
 airflow/decorators/base.py                         |   3 +
 airflow/models/xcom_arg.py                         |   3 +
 .../concepts/dynamic-task-mapping.rst              | 145 ++++++++++++++++++++-
 docs/spelling_wordlist.txt                         |   1 +
 4 files changed, 146 insertions(+), 6 deletions(-)

diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index 448b7e449d..cfc94b525e 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -21,6 +21,7 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
+    ClassVar,
     Collection,
     Dict,
     Generic,
@@ -259,6 +260,8 @@ class _TaskDecorator(Generic[FParams, FReturn, OperatorSubclass]):
 
     decorator_name: str = attr.ib(repr=False, default="task")
 
+    _airflow_is_task_decorator: ClassVar[bool] = True
+
     @multiple_outputs.default
     def _infer_multiple_outputs(self):
         try:
diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py
index a4c2b4d46d..79da1a5466 100644
--- a/airflow/models/xcom_arg.py
+++ b/airflow/models/xcom_arg.py
@@ -358,6 +358,9 @@ class MapXComArg(XComArg):
     """
 
     def __init__(self, arg: XComArg, callables: Sequence[Callable[[Any], Any]]) -> None:
+        for c in callables:
+            if getattr(c, "_airflow_is_task_decorator", False):
+                raise ValueError("map() argument must be a plain function, not a @task operator")
         self.arg = arg
         self.callables = callables
 
diff --git a/docs/apache-airflow/concepts/dynamic-task-mapping.rst b/docs/apache-airflow/concepts/dynamic-task-mapping.rst
index 30622b024c..90db2929ae 100644
--- a/docs/apache-airflow/concepts/dynamic-task-mapping.rst
+++ b/docs/apache-airflow/concepts/dynamic-task-mapping.rst
@@ -190,8 +190,10 @@ If you want to map over the result of a classic operator you will need to create
 
     downstream = MyOperator2.partial(task_id="consumer").expand(input=XComArg(task))
 
-Putting it all together
-=======================
+.. note:: Only a return value can be mapped against. Therefore, it is not allowed to set the ``XComArg``'s ``key`` property manually.
+
+Mixing TaskFlow and classic operators
+=====================================
 
 In this example you have a regular data delivery to an S3 bucket and want to apply the same processing to every file that arrives, no matter how many arrive each time.
 
@@ -206,25 +208,156 @@ In this example you have a regular data delivery to an S3 bucket and want to app
 
 
     with DAG(dag_id="mapped_s3", start_date=datetime(2020, 4, 7)) as dag:
-        files = S3ListOperator(
+        list_filenames = S3ListOperator(
             task_id="get_input",
             bucket="example-bucket",
             prefix='incoming/provider_a/{{ data_interval_start.strftime("%Y-%m-%d") }}',
         )
 
         @task
-        def count_lines(aws_conn_id, bucket, file):
+        def count_lines(aws_conn_id, bucket, filename):
             hook = S3Hook(aws_conn_id=aws_conn_id)
 
-            return len(hook.read_key(file, bucket).splitlines())
+            return len(hook.read_key(filename, bucket).splitlines())
 
         @task
         def total(lines):
             return sum(lines)
 
-        counts = count_lines.partial(aws_conn_id="aws_default", bucket=files.bucket).expand(file=XComArg(files))
+        counts = count_lines.partial(aws_conn_id="aws_default", bucket=list_filenames.bucket).expand(
+            filename=XComArg(list_filenames)
+        )
+
         total(lines=counts)
 
+Assigning multiple parameters to a non-TaskFlow operator
+========================================================
+
+Sometimes an upstream needs to specify multiple arguments to a downstream operator. To do this, you can use the ``expand_kwargs`` function, which takes a sequence of mappings to map against.
+
+.. code-block:: python
+
+    BashOperator.partial(task_id="bash").expand_kwargs(
+        [
+            {"bash_command": "echo $ENV1", "env": {"ENV1": "1"}},
+            {"bash_command": "printf $ENV2", "env": {"ENV2": "2"}},
+        ],
+    )
+
+This produces two task instances at run-time printing ``1`` and ``2`` respectively.
+
+Similar to ``expand``, you can also map against a XCom that returns a list of dicts, or a list of XComs each returning a dict. Re-using the S3 example above, you can use a mapped task to perform "branching" and copy files to different buckets:
+
+.. code-block:: python
+
+    list_filenames = S3ListOperator(...)  # Same as the above example.
+
+
+    @task
+    def create_copy_kwargs(filename):
+        if filename.rsplit(".", 1)[-1] not in ("json", "yml"):
+            dest_bucket_name = "my_text_bucket"
+        else:
+            dest_bucket_name = "my_other_bucket"
+        return {
+            "source_bucket_key": filename,
+            "dest_bucket_key": filename,
+            "dest_bucket_name": dest_bucket_name,
+        }
+
+
+    copy_kwargs = create_copy_kwargs.expand(filename=XComArg(list_filenames))
+
+    # Copy files to another bucket, based on the file's extension.
+    copy_filenames = S3CopyObjectOperator.partial(
+        task_id="copy_files", source_bucket_name=list_filenames.bucket
+    ).expand_kwargs(copy_kwargs)
+
+Filtering items from an expanded task
+=====================================
+
+A mapped task can remove any elements from being passed on to its downstream tasks by returning ``None``. For example, if we want to *only* copy files from an S3 bucket to another with certain extensions, we could implement ``create_copy_kwargs`` like this instead:
+
+.. code-block:: python
+
+    @task
+    def create_copy_kwargs(filename):
+        # Skip files not ending with these suffixes.
+        if filename.rsplit(".", 1)[-1] not in ("json", "yml"):
+            return None
+        return {
+            "source_bucket_key": filename,
+            "dest_bucket_key": filename,
+            "dest_bucket_name": "my_other_bucket",
+        }
+
+
+    # copy_kwargs and copy_files are implemented the same.
+
+This makes ``copy_files`` only expand against ``.json`` and ``.yml`` files, while ignoring the rest.
+
+Transforming mapped data
+========================
+
+Since it is common to want to transform the output data format for task mapping, especially from a non-TaskFlow operator, where the output format is pre-determined and cannot be easily converted (such as ``create_copy_kwargs`` in the above example), a special ``map()`` function can be used to easily perform this kind of transformation. The above example can therefore be modified like this:
+
+.. code-block:: python
+
+    from airflow.exceptions import AirflowSkipException
+
+    filenames = S3ListOperator(...)  # Unchanged.
+
+
+    def create_copy_kwargs(filename):
+        if filename.rsplit(".", 1)[-1] not in ("json", "yml"):
+            raise AirflowSkipException(f"skipping {filename!r}; unexpected suffix")
+        return {
+            "source_bucket_key": filename,
+            "dest_bucket_key": filename,
+            "dest_bucket_name": "my_other_bucket",
+        }
+
+
+    copy_kwargs = XComArg(list_filenames).map(create_copy_kwargs)
+
+    # Unchanged.
+    copy_filenames = S3CopyObjectOperator.partial(...).expand_kwargs(copy_kwargs)
+
+There are a couple of things to note:
+
+#. The callable argument of ``map()`` (``create_copy_kwargs`` in the example) **must not** be a task, but a plain Python function. The transformation is as a part of the "pre-processing" of the downstream task (i.e. ``copy_files``), not a standalone task in the DAG.
+#. The callable always take exactly one positional argument. This function is called for each item in the iterable used for task-mapping, similar to how Python's built-in ``map()`` works.
+#. Since the callable is executed as a part of the downstream task, you can use any existing techniques to write the task function. To mark a component as skipped, for example, you should raise ``AirflowSkipException``. Note that returning ``None`` **does not** work here.
+
+Combining upstream data (aka "zipping")
+=======================================
+
+It is also to want to combine multiple input sources into one task mapping iterable. This is generally known as "zipping" (like Python's built-in ``zip()`` function), and is also performed as pre-processing of the downstream task.
+
+This is especially useful for conditional logic in task mapping. For example, if you want to download files from S3, but rename those files, something like this would be possible:
+
+.. code-block:: python
+
+    list_filenames_a = S3ListOperator(
+        task_id="list_files_in_a",
+        bucket="bucket",
+        prefix="incoming/provider_a/{{ data_interval_start|ds }}",
+    )
+    list_filenames_b = ["rename_1", "rename_2", "rename_3", ...]
+
+    filenames_a_b = XComArg(list_filenames_a).zip(list_filenames_b)
+
+
+    @task
+    def download_filea_from_a_rename(filenames_a_b):
+        fn_a, fn_b = filenames_a_b
+        S3Hook().download_file(fn_a, local_path=fn_b)
+
+
+    download_filea_from_a_rename.expand(filenames_a_b=filenames_a_b)
+
+The ``zip`` function takes arbitrary positional arguments, and return an iterable of tuples of the positional arguments' count. By default, the zipped iterable's length is the same as the shortest of the zipped iterables, with superfluous items dropped. An optional keyword argument ``default`` can be passed to switch the behavior to match Python's ``itertools.zip_longest``—the zipped iterable will have the same length as the *longest* of the zipped iterables, with missing items filled wi [...]
+
 What data types can be expanded?
 ================================
 
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index cb0efadb5c..c54e7ee3fd 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -975,6 +975,7 @@ isfile
 ish
 isn
 iterable
+iterables
 iteratively
 iterm
 itertools