Why branching handles conditional logic in Apache Airflow - See It in Action

Why Branching Handles Conditional Logic in Airflow

📖 Scenario: You are building a simple Airflow workflow to decide which task to run based on a condition. This is like choosing a path in a choose-your-own-adventure story depending on a choice you make.

🎯 Goal: Create an Airflow DAG that uses branching to run different tasks based on a condition. You will first set up the tasks, then add a branching condition, and finally print which task runs.

📋 What You'll Learn

Create three tasks: start, branching, and two possible downstream tasks task_a and task_b.

Use BranchPythonOperator to decide which task to run based on a variable run_task_a.

Print the name of the task that runs after branching.

💡 Why This Matters

🌍 Real World

Branching in Airflow is used to run different tasks based on conditions like success/failure, data values, or external triggers.

💼 Career

Understanding branching helps you build flexible workflows that adapt to different situations, a key skill for workflow automation and data engineering roles.

1

Create initial Airflow tasks

Create an Airflow DAG named branching_dag with a start task using DummyOperator and two tasks task_a and task_b also using DummyOperator.

Apache Airflow

# Your code here to create DAG and three DummyOperator tasks: start, task_a, task_b

Hint

Use DummyOperator to create simple placeholder tasks.

2

Add branching condition variable

Create a variable called run_task_a and set it to True. This variable will decide which task to run next.

Apache Airflow

from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime

default_args = {'start_date': datetime(2024, 1, 1)}

dag = DAG('branching_dag', default_args=default_args, schedule_interval=None)

start = DummyOperator(task_id='start', dag=dag)
task_a = DummyOperator(task_id='task_a', dag=dag)
task_b = DummyOperator(task_id='task_b', dag=dag)

# Create a variable run_task_a and set it to True
# Your code here

Hint

This variable will be used later to decide which task to run.

3

Add BranchPythonOperator with condition

Create a function called choose_task that returns 'task_a' if run_task_a is True, else returns 'task_b'. Then create a BranchPythonOperator named branching that uses this function. Set the task dependencies so start runs before branching, and branching runs before both task_a and task_b.

Apache Airflow

from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime
from airflow.operators.python import BranchPythonOperator

default_args = {'start_date': datetime(2024, 1, 1)}

dag = DAG('branching_dag', default_args=default_args, schedule_interval=None)

start = DummyOperator(task_id='start', dag=dag)
task_a = DummyOperator(task_id='task_a', dag=dag)
task_b = DummyOperator(task_id='task_b', dag=dag)

run_task_a = True

# Define choose_task function and BranchPythonOperator branching
# Set dependencies: start >> branching >> [task_a, task_b]
# Your code here

Hint

Use BranchPythonOperator to decide which task runs next based on the function's return value.

4

Print which task runs after branching

Add a print statement that outputs the result of calling choose_task(). This shows which task will run after the branching decision.

Apache Airflow

from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime
from airflow.operators.python import BranchPythonOperator

default_args = {'start_date': datetime(2024, 1, 1)}

dag = DAG('branching_dag', default_args=default_args, schedule_interval=None)

start = DummyOperator(task_id='start', dag=dag)
task_a = DummyOperator(task_id='task_a', dag=dag)
task_b = DummyOperator(task_id='task_b', dag=dag)

run_task_a = True

def choose_task():
    if run_task_a:
        return 'task_a'
    else:
        return 'task_b'

branching = BranchPythonOperator(
    task_id='branching',
    python_callable=choose_task,
    dag=dag
)

start >> branching >> [task_a, task_b]

# Print which task will run after branching
# Your code here

Hint

Calling choose_task() returns the task id that will run next.

Practice

1. What is the main purpose of branching in Airflow workflows?

easy

A. To choose which tasks run based on conditions

B. To run all tasks in parallel

C. To schedule tasks at fixed intervals

D. To retry failed tasks automatically

Solution

1. Step 1: Understand branching concept

   Branching allows workflows to decide which path to take based on conditions.

2. Step 2: Compare with other options

   Running all tasks in parallel or scheduling are different features, not branching.

3. Final Answer:

   To choose which tasks run based on conditions -> Option A

4. Quick Check:

   Branching = Conditional path selection [OK]

Hint: Branching picks one path based on condition [OK]

Common Mistakes:

• Confusing branching with parallel execution
• Thinking branching schedules tasks
• Assuming branching retries tasks

2. Which Airflow operator is used to implement branching logic?

easy

A. PythonOperator

B. DummyOperator

C. BashOperator

D. BranchPythonOperator

Solution

1. Step 1: Identify branching operator

   The BranchPythonOperator is designed to choose the next task based on a Python function's return value.

2. Step 2: Eliminate other operators

   PythonOperator runs Python code but does not branch; DummyOperator is a placeholder; BashOperator runs shell commands.

3. Final Answer:

   BranchPythonOperator -> Option D

4. Quick Check:

   Branching operator = BranchPythonOperator [OK]

Hint: BranchPythonOperator controls conditional task flow [OK]

Common Mistakes:

• Choosing PythonOperator instead of BranchPythonOperator
• Confusing DummyOperator as branching
• Using BashOperator for branching

3. Given this BranchPythonOperator function, which task will run next?

def choose_task():
    if 5 > 3:
        return 'task_a'
    else:
        return 'task_b'

medium

A. task_b

B. Both task_a and task_b

C. task_a

D. No task runs

Solution

1. Step 1: Evaluate the condition in the function

   5 > 3 is True, so the function returns 'task_a'.

2. Step 2: Understand branching behavior

   Only the task returned by the function runs; others are skipped.

3. Final Answer:

   task_a -> Option C

4. Quick Check:

   Condition true -> task_a runs [OK]

Hint: Return task name to run next [OK]

Common Mistakes:

• Thinking both tasks run
• Choosing task_b despite condition
• Assuming no task runs

4. What is wrong with this BranchPythonOperator function?

def choose_task():
    if x > 10:
        return 'task_x'
    else:
        return 'task_y'

medium

A. Variable x is not defined, causing an error

B. The function does not return a task id

C. BranchPythonOperator cannot return strings

D. The else block is missing

Solution

1. Step 1: Check variable usage

   The variable x is used but not defined anywhere in the function or globally.

2. Step 2: Understand error cause

   Using an undefined variable causes a NameError at runtime.

3. Final Answer:

   Variable x is not defined, causing an error -> Option A

4. Quick Check:

   Undefined variable = runtime error [OK]

Hint: Check all variables are defined before use [OK]

Common Mistakes:

• Assuming missing else block causes error
• Thinking BranchPythonOperator can't return strings
• Ignoring undefined variable errors

5. You want to run task_a if a file exists, otherwise run task_b. Which BranchPythonOperator function correctly implements this?

hard

A. def choose():
if os.path.exists('/path/file'):
print('task_a')
else:
print('task_b')

B. def choose():
if os.path.exists('/path/file'):
return 'task_a'
return 'task_b'

C. def choose():
if os.path.exists('/path/file'):
return ['task_a', 'task_b']
else:
return 'task_b'

D. def choose():
if os.path.exists('/path/file'):
return 'task_b'
else:
return 'task_a'

Solution

1. Step 1: Check condition and return values

   def choose(): if os.path.exists('/path/file'): return 'task_a' return 'task_b' returns 'task_a' if file exists, else 'task_b', matching the requirement.

2. Step 2: Identify errors in other options

   def choose(): if os.path.exists('/path/file'): return 'task_b' else: return 'task_a' reverses tasks; def choose(): if os.path.exists('/path/file'): return ['task_a', 'task_b'] else: return 'task_b' returns a list which is invalid; def choose(): if os.path.exists('/path/file'): print('task_a') else: print('task_b') prints instead of returning task id.

3. Final Answer:

   def choose(): if os.path.exists('/path/file'): return 'task_a' return 'task_b' -> Option B

4. Quick Check:

   Return correct task id based on condition [OK]

Hint: Return one task id string, not list or print [OK]

Common Mistakes:

• Returning list of tasks instead of single task
• Printing task names instead of returning
• Swapping task names in condition