TaskFlow API for cleaner XCom in Apache Airflow - Step-by-Step Execution

Process Flow - TaskFlow API for cleaner XCom

Define Python function task

↓

Decorate with @task

↓

Use function in DAG context

↓

Return value auto-pushed to XCom

↓

Downstream task receives value as argument

↓

No manual XCom push/pull needed

↓

DAG runs cleanly

The TaskFlow API lets you write Python functions as tasks that automatically pass data via XCom without manual push or pull.

Execution Sample

Apache Airflow

from airflow.decorators import dag, task
from datetime import datetime

@dag(start_date=datetime(2024,1,1), schedule_interval='@daily', catchup=False)
def my_dag():
    @task
    def get_data():
        return 'hello'

    @task
    def print_data(data):
        print(data)

    print_data(get_data())

dag = my_dag()

This DAG defines two tasks using TaskFlow API; get_data returns a string, print_data prints it, passing data automatically via XCom.

Process Table

Step	Action	Task	XCom Push	XCom Pull	Output/Result
1	Run get_data task	get_data	'hello' pushed automatically	None	Returns 'hello'
2	Run print_data task with argument	print_data	None	'hello' pulled automatically	Prints 'hello'
3	DAG completes	all	All XComs handled by TaskFlow	All XComs handled by TaskFlow	Clean DAG run with no manual XCom code

💡 All tasks complete; data passed automatically via TaskFlow API without manual XCom push/pull.

Status Tracker

Variable	Start	After get_data	After print_data	Final
data	None	'hello' returned by get_data	'hello' received by print_data	'hello' printed

Key Moments - 2 Insights

Why don't we see any manual XCom push or pull commands in the code?

Because the TaskFlow API automatically pushes the return value of a @task function to XCom and pulls it when passed as an argument, as shown in execution_table steps 1 and 2.

How does print_data get the value from get_data without explicit XCom calls?

The TaskFlow API handles this by automatically pulling the XCom value returned by get_data when print_data is called with get_data() as argument, as seen in execution_table step 2.

Visual Quiz - 3 Questions

Look at the execution table, what does the get_data task push to XCom?

A None

B 'hello'

C print_data function

D A datetime object

💡 Hint

Check the 'XCom Push' column in row 1 of the execution_table.

At which step does print_data pull data from XCom automatically?

A Step 2

B Step 1

C Step 3

D Never

💡 Hint

Look at the 'XCom Pull' column in the execution_table for step 2.

If we changed get_data to return a number instead of 'hello', what would print_data receive?

A The string 'hello'

B None

C The number returned by get_data

D An error

💡 Hint

Refer to variable_tracker showing how data flows from get_data to print_data.

Concept Snapshot

TaskFlow API uses @task decorator to turn Python functions into Airflow tasks.
Return values are automatically pushed to XCom.
Passing a task's output as argument to another task auto-pulls XCom.
No manual XCom push/pull needed.
Simplifies data passing and makes DAG code cleaner.

Full Transcript

The TaskFlow API in Airflow lets you write Python functions as tasks using the @task decorator. When a task function returns a value, Airflow automatically pushes it to XCom. When another task function takes that task's output as an argument, Airflow automatically pulls the value from XCom. This removes the need for manual XCom push and pull commands. In the example DAG, get_data returns 'hello' which is automatically pushed to XCom. The print_data task receives this value as an argument, automatically pulling it from XCom, and prints it. This makes DAGs cleaner and easier to read.

Practice

1. What is the main benefit of using the TaskFlow API in Airflow for passing data between tasks?

easy

A. It automatically handles data sharing between tasks using return values.

B. It requires manual XCom push and pull commands for data sharing.

C. It disables task dependencies to speed up execution.

D. It only works with BashOperator tasks.

Solution

1. Step 1: Understand TaskFlow API purpose

   The TaskFlow API simplifies task communication by using Python functions and decorators.

2. Step 2: Recognize automatic XCom handling

   Return values from decorated tasks are automatically shared with downstream tasks via XCom without manual commands.

3. Final Answer:

   It automatically handles data sharing between tasks using return values. -> Option A

4. Quick Check:

   TaskFlow API = automatic XCom via return [OK]

Hint: TaskFlow API shares data automatically via return values [OK]

Common Mistakes:

• Thinking manual XCom push/pull is needed
• Believing it disables dependencies
• Assuming it only works with BashOperator

2. Which of the following is the correct way to define a task using the TaskFlow API in Airflow?

easy

A. def my_task(): pass

B. @dag
def my_task(): pass

C. @task
def my_task(): pass

D. my_task = PythonOperator(task_id='my_task')

Solution

1. Step 1: Identify TaskFlow API syntax

   Tasks are defined as Python functions decorated with @task to enable TaskFlow features.

2. Step 2: Compare options

   @task def my_task(): pass uses @task decorator correctly; others either lack decorator or use wrong decorators/operators.

3. Final Answer:

   @task\ndef my_task(): pass -> Option C

4. Quick Check:

   TaskFlow task = function + @task decorator [OK]

Hint: Use @task decorator above function to define TaskFlow task [OK]

Common Mistakes:

• Omitting @task decorator
• Confusing @dag with @task
• Using PythonOperator instead of TaskFlow

3. Given the following code snippet using TaskFlow API, what will be printed?

@task
def add(x, y):
    return x + y

@task
def print_result(value):
    print(f'Result is {value}')

with DAG('example_dag', start_date=datetime(2023,1,1)) as dag:
    sum_value = add(3, 4)
    print_result(sum_value)

medium

A. No output printed

B. Result is

C. SyntaxError

D. Result is 7

Solution

1. Step 1: Understand TaskFlow return and XCom

   The add task returns 7, which is passed automatically to print_result.

2. Step 2: Trace print statement execution

   print_result prints the string with the value 7, so output is 'Result is 7'.

3. Final Answer:

   Result is 7 -> Option D

4. Quick Check:

   Returned value passed and printed = 'Result is 7' [OK]

Hint: Returned values flow automatically and print as expected [OK]

Common Mistakes:

• Expecting raw XCom object printed
• Confusing syntax error with correct code
• Thinking no output prints without manual XCom pull

4. What is wrong with this TaskFlow API code snippet?

@task
def multiply(x, y):
    print(x * y)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
    result = multiply(2, 5)
    print(f'Result is {result}')

medium

A. The DAG is missing the schedule_interval parameter.

B. The multiply task should return the result instead of printing it.

C. The print statement inside the DAG is invalid syntax.

D. The @task decorator is missing parentheses.

Solution

1. Step 1: Identify how TaskFlow passes data

   TaskFlow passes data via return values, not print statements inside tasks.

2. Step 2: Analyze code behavior

   Since multiply only prints and returns None, result is None, so printing it shows 'Result is None'.

3. Final Answer:

   The multiply task should return the result instead of printing it. -> Option B

4. Quick Check:

   Return value needed for XCom, not print [OK]

Hint: Return values to share data, not print inside tasks [OK]

Common Mistakes:

• Thinking print shares data between tasks
• Confusing DAG parameters with task errors
• Believing @task needs parentheses always

5. You want to create a DAG using TaskFlow API where task A returns a list of numbers, and task B receives this list and returns the sum. Which code snippet correctly implements this pattern?

hard

A. @task
def task_a():
return [1, 2, 3]

@task
def task_b(numbers):
return sum(numbers)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
nums = task_a()
total = task_b(nums)

B. @task
def task_a():
print([1, 2, 3])

@task
def task_b(numbers):
print(sum(numbers))

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
nums = task_a()
total = task_b(nums)

C. @task
def task_a():
return '1,2,3'

@task
def task_b(numbers):
return sum(numbers)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
nums = task_a()
total = task_b(nums)

D. @task
def task_a():
return [1, 2, 3]

@task
def task_b(numbers):
return sum(numbers)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
total = task_b(task_a())

Solution

1. Step 1: Check task A returns a list

   @task
   def task_a():
       return [1, 2, 3]
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   returns a list [1, 2, 3] correctly;

   @task
   def task_a():
       print([1, 2, 3])
   
   @task
   def task_b(numbers):
       print(sum(numbers))
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   prints instead of returning;

   @task
   def task_a():
       return '1,2,3'
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   returns a string, not list.

2. Step 2: Check task B sums the list argument

   @task
   def task_a():
       return [1, 2, 3]
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   passes the list to task B which sums it correctly;

   @task
   def task_a():
       return [1, 2, 3]
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       total = task_b(task_a())

   skips assigning intermediate variable but is valid but less clear.

3. Final Answer:

   Option A correctly returns a list and passes it to sum in task B. -> Option A

4. Quick Check:

   Return list and pass it to sum task [OK]

Hint: Return data structures, pass as args, sum in next task [OK]

Common Mistakes:

• Using print instead of return in tasks
• Passing string instead of list to sum
• Ignoring task dependencies or variable assignment