TaskFlow API for cleaner XCom in Apache Airflow - Practice Problems & Coding Challenges

Challenge - 5 Problems

💻 Command Output

intermediate

What is the output of this Airflow TaskFlow function?

Given the following Airflow DAG task using TaskFlow API, what will be the output when the task runs successfully?

Apache Airflow

from airflow.decorators import task

@task
def multiply_by_two(x: int) -> int:
    return x * 2

result = multiply_by_two(5)
print(result)

A XComArg object

B 10

C None

D TypeError

💡 Hint

Remember that TaskFlow API returns a special object representing the task output, not the raw value.

In Airflow's TaskFlow API, calling a decorated task function returns an XComArg object, which represents the future result of the task. The actual value is pushed to XCom only after the task runs.

🧠 Conceptual

intermediate

Which statement best describes XCom behavior in TaskFlow API?

Select the correct statement about how XComs are handled when using Airflow's TaskFlow API.

A TaskFlow API disables XComs by default to improve performance.

B You must manually push XCom values using task_instance.xcom_push() inside the task.

C TaskFlow API automatically serializes and pushes return values to XCom without extra code.

D XComs are only available if you use PythonOperator, not TaskFlow API.

💡 Hint

Think about how TaskFlow API simplifies passing data between tasks.

TaskFlow API automatically handles pushing the return value of a decorated task to XCom, so you don't need to manually push or pull values.

❓ Troubleshoot

advanced

Why does this TaskFlow task fail to push the expected XCom value?

Consider this task code snippet: from airflow.decorators import task @task def add_numbers(a, b): sum = a + b What is the reason the XCom value is None after task execution?

A The function does not return the sum, so nothing is pushed to XCom.

B The @task decorator requires explicit xcom_push calls inside the function.

C The variables a and b are not typed, causing serialization failure.

D Airflow disables XCom for functions with more than one argument.

💡 Hint

Check if the function returns a value to be pushed to XCom.

In TaskFlow API, only the return value of the decorated function is pushed to XCom. Without a return statement, the XCom value is None.

🔀 Workflow

advanced

What is the correct way to pass data between two TaskFlow tasks?

Given two tasks defined with TaskFlow API: @task def task_a(): return 42 @task def task_b(x): return x + 1 How do you correctly set up the DAG to pass the output of task_a to task_b?

A

result = task_a()
task_b()  # no argument

B

task_a()
task_b(task_a())  # call twice

C final = task_b(task_a) # pass function, not result

D

result = task_a()
final = task_b(result)

💡 Hint

Remember that calling a task returns an XComArg that can be passed as argument to another task.

Calling task_a() returns an XComArg object representing its output, which can be passed directly as argument to task_b().

✅ Best Practice

expert

Which approach best improves XCom data handling in TaskFlow API for large data?

You have a TaskFlow task that returns a large dictionary. Which method is best to avoid performance issues with XCom?

A Serialize the dictionary as a JSON string and return it directly from the task.

B Store the large data in external storage (e.g., S3) and push only the reference (path) via XCom.

C Split the dictionary into smaller chunks and return multiple tasks each pushing part of data.

D Increase Airflow's XCom max size limit to accommodate the large dictionary.

💡 Hint

Think about how to keep XCom lightweight and efficient.

XCom is designed for small messages. For large data, best practice is to store it externally and pass only a reference via XCom to avoid performance degradation.

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