TaskFlow API for cleaner XCom in Apache Airflow - Mini Project: Build & Apply

TaskFlow API for Cleaner XCom in Airflow

📖 Scenario: You are working with Apache Airflow to automate a simple data processing workflow. You want to use the TaskFlow API to pass data between tasks cleanly without manually handling XComs.

🎯 Goal: Build a simple Airflow DAG using the TaskFlow API where one task returns a value and the next task receives it as input automatically.

📋 What You'll Learn

Create a DAG with two tasks using the TaskFlow API

First task returns a string value

Second task receives the returned value as a parameter

Print the received value in the second task

💡 Why This Matters

🌍 Real World

In real projects, Airflow workflows often need to pass data between tasks. The TaskFlow API makes this easier and cleaner by using Python functions and automatic XCom handling.

💼 Career

Understanding the TaskFlow API is important for data engineers and DevOps professionals who build and maintain reliable, readable, and maintainable data pipelines using Airflow.

1

Create the DAG and first task

Create a DAG called taskflow_dag with start_date set to days_ago(1) and schedule_interval set to None. Then create a Python function called extract_data that returns the string 'Hello from TaskFlow'. Use the @task decorator from airflow.decorators on this function.

Apache Airflow

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.decorators import task

# Create DAG and define extract_data task
# Your code here

Hint

Use with DAG(...): to create the DAG context. Decorate the function with @task and return the exact string.

2

Create the second task to receive data

Inside the same DAG context, create a Python function called print_data decorated with @task that takes one parameter called data. Inside the function, write a print statement to display the data parameter.

Apache Airflow

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.decorators import task

with DAG('taskflow_dag', start_date=days_ago(1), schedule_interval=None) as dag:
    @task
    def extract_data():
        return 'Hello from TaskFlow'

    # Create print_data task that accepts data parameter
    # Your code here

Hint

Define print_data with one parameter data and print it inside the function.

3

Link the tasks using TaskFlow API

Call the extract_data task and assign its result to a variable called extracted. Then call the print_data task passing extracted as the argument. This will link the tasks and pass data automatically.

Apache Airflow

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.decorators import task

with DAG('taskflow_dag', start_date=days_ago(1), schedule_interval=None) as dag:
    @task
    def extract_data():
        return 'Hello from TaskFlow'

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

    # Call extract_data and assign to extracted
    # Call print_data with extracted
    # Your code here

Hint

Assign the result of extract_data() to extracted and pass it to print_data(extracted).

4

Run the DAG and check output

Run the DAG and observe the output of the print_data task. The output should be exactly Hello from TaskFlow printed in the task logs.

Apache Airflow

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.decorators import task

with DAG('taskflow_dag', start_date=days_ago(1), schedule_interval=None) as dag:
    @task
    def extract_data():
        return 'Hello from TaskFlow'

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

    extracted = extract_data()
    print_data(extracted)

# Run the DAG and check output

Hint

Check the logs of the print_data task after running the DAG to see the printed message.

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