Introduction to the TaskFlow API and Airflow decorators

The TaskFlow API is a functional API for using decorators to define DAGs and tasks, which simplifies the process for passing data between tasks and defining dependencies. You can use TaskFlow decorator functions (for example, @task) to pass data between tasks by providing the output of one task as an argument to another task. Decorators are a simpler, cleaner way to define your tasks and DAGs and can be used in combination with traditional operators.

In this guide, you’ll learn about the benefits of decorators and the decorators available in Airflow. You’ll also review an example DAG and learn when you should use decorators and how you can combine them with traditional operators in a DAG.

Assumed knowledge

To get the most out of this guide, you should have an understanding of:

• Basic Python. See the Python Documentation.
• Airflow operators. See Operators 101.

What is a decorator?

In Python, decorators are functions that take another function as an argument and extend the behavior of that function. For example, the @multiply_by_100_decorator takes any function as the decorated_function argument and returns the result of that function multiplied by 100.

1
# definition of the decorator function
2
def multiply_by_100_decorator(decorated_function):
3
    def wrapper(num1, num2):
4
        result = decorated_function(num1, num2) * 100
5
        return result
6
7
    return wrapper
8
9
# definition of the `add` function decorated with the `multiply_by_100_decorator`
10
@multiply_by_100_decorator
11
def add(num1, num2):
12
    return num1 + num2
13
14
# definition of the `subtract` function decorated with the `multiply_by_100_decorator`
15
@multiply_by_100_decorator
16
def subtract(num1, num2):
17
    return num1 - num2
18
19
# calling the decorated functions
20
print(add(1, 9))  # prints 1000
21
print(subtract(4, 2))  # prints 200

In the context of Airflow, decorators contain more functionality than this simple example, but the basic idea is the same: the Airflow decorator function extends the behavior of a normal Python function to turn it into an Airflow task, task group or DAG.

When to use the TaskFlow API

The purpose of the TaskFlow API in Airflow is to simplify the DAG authoring experience by eliminating the boilerplate code required by traditional operators. The result can be cleaner DAG files that are more concise and easier to read.

In general, whether you use the TaskFlow API is a matter of your own preference and style. In most cases, a TaskFlow decorator and the corresponding traditional operator will have the same functionality. You can also mix decorators and traditional operators within a single DAG.

How to use the TaskFlow API

The TaskFlow API allows you to write your Python tasks with decorators. It handles passing data between tasks using XCom and infers task dependencies automatically.

Using decorators to define your Python functions as tasks is easy. Let’s take a before and after example. Under the Traditional syntax tab below, there is a basic ETL DAG with tasks to get data from an API, process the data, and store it. Click on the Decorators tab to see the same DAG written using Airflow decorators.

1
import logging
2
from datetime import datetime
3
4
import requests
5
from airflow import DAG
6
from airflow.operators.python import PythonOperator
7
8
API = "https://api.coingecko.com/api/v3/simple/price?ids=bitcoin&vs_currencies=usd&include_market_cap=true&include_24hr_vol=true&include_24hr_change=true&include_last_updated_at=true"
9
10
11
def _extract_bitcoin_price():
12
    return requests.get(API).json()["bitcoin"]
13
14
15
def _process_data(ti):
16
    response = ti.xcom_pull(task_ids="extract_bitcoin_price")
17
    logging.info(response)
18
    processed_data = {"usd": response["usd"], "change": response["usd_24h_change"]}
19
    ti.xcom_push(key="processed_data", value=processed_data)
20
21
22
def _store_data(ti):
23
    data = ti.xcom_pull(task_ids="process_data", key="processed_data")
24
    logging.info(f"Store: {data['usd']} with change {data['change']}")
25
26
27
with DAG(
28
    "classic_dag", schedule="@daily", start_date=datetime(2021, 12, 1), catchup=False
29
):
30
    extract_bitcoin_price = PythonOperator(
31
        task_id="extract_bitcoin_price", python_callable=_extract_bitcoin_price
32
    )
33
34
    process_data = PythonOperator(task_id="process_data", python_callable=_process_data)
35
36
    store_data = PythonOperator(task_id="store_data", python_callable=_store_data)
37
38
    extract_bitcoin_price >> process_data >> store_data

1
import logging
2
from datetime import datetime
3
from typing import Dict
4
5
import requests
6
from airflow.decorators import dag, task
7
8
API = "https://api.coingecko.com/api/v3/simple/price?ids=bitcoin&vs_currencies=usd&include_market_cap=true&include_24hr_vol=true&include_24hr_change=true&include_last_updated_at=true"
9
10
11
@dag(schedule="@daily", start_date=datetime(2021, 12, 1), catchup=False)
12
def taskflow():
13
    @task(task_id="extract", retries=2)
14
    def extract_bitcoin_price() -> Dict[str, float]:
15
        return requests.get(API).json()["bitcoin"]
16
17
    @task(multiple_outputs=True)
18
    def process_data(response: Dict[str, float]) -> Dict[str, float]:
19
        logging.info(response)
20
        return {"usd": response["usd"], "change": response["usd_24h_change"]}
21
22
    @task
23
    def store_data(data: Dict[str, float]):
24
        logging.info(f"Store: {data['usd']} with change {data['change']}")
25
26
    store_data(process_data(extract_bitcoin_price()))
27
28
29
taskflow()

The decorated version of the DAG eliminates the need to explicitly instantiate the PythonOperator, has much less code and is easier to read. Notice that it also doesn’t require using ti.xcom_pull and ti.xcom_push to pass data between tasks. This is all handled by the TaskFlow API when you define your task dependencies with store_data(process_data(extract_bitcoin_price())).

Here are some other things to keep in mind when using decorators:

• You must call all decorated functions in your DAG file so that Airflow can register the task or DAG. For example, taskflow() is called at the end of the previous example to call the DAG function.

• When you define a task, the task_id defaults to the name of the function you decorated. If you want to change this behavior, you can pass a task_id to the decorator as shown in the extract task example. Similarly, other BaseOperator task-level parameters, such as retries or pool, can be defined within the decorator:

  1
  @task(
  2
      task_id="say_hello_world"
  3
      retries=3,
  4
      pool="my_pool",
  5
  )
  6
  def taskflow_func():
  7
      retun "Hello World"
  8
  9
  taskflow_func()  # this creates a task with the task_id `say_hello_world`

• Override task-level parameters when you call the task by using the .override() method. The () at the end of the line calls the task with the overridden parameters applied.

  1
  # this creates a task with the task_id `greeting`
  2
  taskflow_func.override(retries=5, pool="my_other_pool", task_id="greeting")()

• If you call the same task multiple times and do not override the task_id, Airflow creates multiple unique task IDs by appending a number to the end of the original task ID (for example, say_hello, say_hello__1, say_hello__2, etc). You can see the result of this in the following example:

  1
  # task definition
  2
  @task
  3
  def say_hello(dog):
  4
      return f"Hello {dog}!"
  5
  6
  ### calling the task 4 times, creating 4 tasks in the DAG
  7
  # this task will have the id `say_hello` and print "Hello Avery!"
  8
  say_hello("Avery")
  9
  # this task will have the id `greet_dog` and print "Hello Piglet!"
  10
  say_hello.override(task_id="greet_dog")("Piglet")
  11
  # this task will have the id `say_hello__1` and print "Hello Peanut!"
  12
  say_hello("Peanut")
  13
  # this task will have the id `say_hello__2` and print "Hello Butter!"
  14
  say_hello("Butter")

• You can decorate a function that is imported from another file as shown in the following code snippet:

  1
  from include.my_file import my_function
  2
  3
  @task
  4
  def taskflow_func():
  5
      my_function()

  This is recommended in cases where you have lengthy Python functions since it will make your DAG file easier to read.

• You can assign the output of a called decorated task to a Python object to be passed as an argument into another decorated task. This is helpful when the output of one decorated task is needed in several downstream functions.

  1
  @task 
  2
  def get_fruit_options():
  3
      return ["peach", "raspberry", "pineapple"]
  4
  5
  @task
  6
  def eat_a_fruit(list):
  7
      index = random.randint(0, len(list) - 1)
  8
      print(f"I'm eating a {list[index]}!")
  9
  10
  @task 
  11
  def gift_a_fruit(list):
  12
      index = random.randint(0, len(list) - 1)
  13
      print(f"I'm giving you a {list[index]}!")
  14
  15
  # you can assign the output of a decorated task to a Python object
  16
  my_fruits = get_fruit_options()
  17
  eat_a_fruit(my_fruits)
  18
  gift_a_fruit(my_fruits)

View more examples on how to use Airflow task decorators in the Astronomer webinars and the Apache Airflow TaskFlow API tutorial.

Mixing TaskFlow decorators with traditional operators

If you have a DAG that uses PythonOperator and other operators that don’t have decorators, you can easily combine decorated functions and traditional operators in the same DAG. For example, you can add an EmailOperator to the previous example by updating your code to the following:

1
import logging
2
from datetime import datetime
3
from typing import Dict
4
5
import requests
6
from airflow.decorators import dag, task
7
from airflow.operators.email import EmailOperator
8
9
API = "https://api.coingecko.com/api/v3/simple/price?ids=bitcoin&vs_currencies=usd&include_market_cap=true&include_24hr_vol=true&include_24hr_change=true&include_last_updated_at=true"
10
11
12
@dag(schedule="@daily", start_date=datetime(2021, 12, 1), catchup=False)
13
def taskflow():
14
    @task(task_id="extract", retries=2)
15
    def extract_bitcoin_price() -> Dict[str, float]:
16
        return requests.get(API).json()["bitcoin"]
17
18
    @task(multiple_outputs=True)
19
    def process_data(response: Dict[str, float]) -> Dict[str, float]:
20
        logging.info(response)
21
        return {"usd": response["usd"], "change": response["usd_24h_change"]}
22
23
    @task
24
    def store_data(data: Dict[str, float]):
25
        logging.info(f"Store: {data['usd']} with change {data['change']}")
26
27
    email_notification = EmailOperator(
28
        task_id="email_notification",
29
        to="noreply@astronomer.io",
30
        subject="dag completed",
31
        html_content="the dag has finished",
32
    )
33
34
    store_data(process_data(extract_bitcoin_price())) >> email_notification
35
36
37
taskflow()

Note that when adding traditional operators, dependencies are still defined using bit-shift operators.

You can pass information between decorated tasks and traditional operators using XCom. See the following sections for examples.

TaskFlow to TaskFlow

If both tasks are defined using the TaskFlow API, you can pass information directly between them by providing the called task function as a positional argument to the downstream task. Airflow will infer the dependency between the two tasks.

1
@task
2
def get_23_TF():
3
    return 23
4
5
@task
6
def plus_10_TF(x):
7
    return x + 10
8
9
plus_10_TF(get_23_TF())  # plus_10_TF will return 33
10
# or `plus_10_TF(x=get_23_TF())` if you want to use kwargs

Traditional operator to TaskFlow

Pass the .output values of traditional Operator-based tasks to the callable of TaskFlow tasks to automatically create a relationship between the two tasks. For example:

1
def first_task_callable():
2
    return "hello world"
3
4
first_task = PythonOperator(
5
    task_id="first_task",
6
    python_callable=first_task_callable
7
)
8
9
first_task_result = first_task.output
10
11
@task
12
def second_task(first_task_result_value):
13
    return f"{first_task_result_value} and hello again"
14
15
# Providing first_task_result (traditional operator result) as an argument to a second_task
16
# (a TaskFlow function) automatically registers the dependency between the tasks.
17
second_task(first_task_result)

When no outputs are passed as arguments to automatically register dependencies, but you still want one task to complete before the second starts, use >> to create the relationship between the first task and the second task:

1
@task
2
def first_task():
3
    return "hello world"
4
5
second_task = EmptyOperator(
6
    task_id="second_task",
7
)
8
9
first_task() >> second_task

The data a task returns varies based on its operator. A searchable registry of operators at the Airflow Registry describes what operators return and documents any operator parameters that can be used to control the return format.

TaskFlow to traditional operator

TaskFlow tasks, when called, return a reference (referred to internally as an XComArg) that can be passed into templateable fields of traditional operators to automatically create a relationship between the two tasks.

The list of templateable fields varies by operator. A searchable registry of operators at the Airflow Registry details which fields are templateable by default and users can control which fields are templateable.

Here’s how you could pass the result of a TaskFlow function to a traditional PythonOperator’s callable via an argument:

1
@task
2
def first_task():
3
    return "hello"
4
5
first_task_result = first_task()
6
7
def second_task_callable(x):
8
    uppercase_text = x.upper()
9
    return f"Upper cased version of prior task result: {uppercase_text}"
10
11
# when first_task_result (an XComArg) is provided as an argument to a templated function (op_args)
12
# Airflow automatically registers that second_task depends on first_task
13
second_task = PythonOperator(
14
    task_id="second_task",
15
    python_callable=second_task_callable,
16
    op_args=[ first_task_result ] # note op_args requires a LIST of XComArgs
17
)
18
19
# Astronomer recommends against explicitly defining this redundant-reference
20
# first_task_result >> second_task

When only the order of task execution is important, don’t pass the return value of the first task as a parameter to the second task - instead use >> to explicitly create the relationship. For example:

1
@task
2
def first_task():
3
    return "hello world"
4
5
second_task = EmptyOperator(task_id="second_task")
6
7
first_task() >> second_task

Traditional operator to traditional operator

For the sake of completeness the below example shows how to use the output of one traditional operator in another traditional operator by accessing the .output attribute of the upstream task. The dependency has to be defined explicitly using bit-shift operators.

1
def get_23_traditional():
2
    return 23
3
4
def plus_10_traditional(x):
5
    return x + 10
6
7
get_23_task = PythonOperator(
8
    task_id="get_23_task",
9
    python_callable=get_23_traditional
10
)
11
12
plus_10_task = PythonOperator(
13
    task_id="plus_10_task",
14
    python_callable=plus_10_traditional,
15
    op_args=[get_23_task.output]
16
)
17
18
# plus_10_task will return 33
19
20
# when only using traditional operators, define dependencies explicitly
21
get_23_task >> plus_10_task

Available Airflow decorators

There are several decorators available to use with Airflow. This list provides a reference of currently available decorators:

• DAG decorator (@dag()), which creates a DAG.
• TaskGroup decorator (@task_group()), which creates a TaskGroup.
• Task decorator (@task()), which creates a Python task.
• Bash decorator (@task.bash()) which creates a BashOperator task.
• Python Virtual Env decorator (@task.virtualenv()), which runs your Python task in a virtual environment.
• Docker decorator (@task.docker()), which creates a DockerOperator task.
• Short circuit decorator (@task.short_circuit()), which evaluates a condition and skips downstream tasks if the condition is False.
• Branch decorator (@task.branch()), which creates a branch in your DAG based on an evaluated condition.
• BranchExternalPython decorator (@task.branch_external_python), which creates a branch in your DAG running Python code in a pre-existing virtual environment.
• BranchPythonVirtualenvOperator (@task.branch_virtualenv), which creates a branch in your DAG running Python code in a newly created virtual environment. The environment can be cached by providing a venv_cache_path.
• Kubernetes pod decorator (@task.kubernetes()), which runs a KubernetesPodOperator task.
• Sensor decorator (@task.sensor()), which turns a Python function into a sensor.
• PySpark decorator (@task.pyspark()), which is injected with a SparkSession and SparkContext object if available.

You can also create your own custom task decorator.