[Airflow] BranchPython 오퍼레이터로 분기처리하기

 

1. Task 분기처리 유형

• Task 분기처리가 필요한 이유
  task1의 결과에 따라 task2-x 중 하나만 수행하도록 구성해야 할 때 필요함
  ex.
  'Good', 'Bad', 'Pending' 이라는 결과 3개 중 하나가 나오고
  그에 따라 task2-1~3 중 하나가 실행되도록 해야할 경우

 

 

• Task 분기처리 방법
  1) BranchPythonOperator (주로사용)
  2) task.branch 데코레이터 이용 (주로사용)
  3) BaseBranchOperator 상속하여 직접 개발
• 공통적으로 리턴 값으로 후행 task의 id를 str 또는 list로 리턴해야 함
• 3가지 분기처리 방법은 방법만 다를 뿐 결과는 동일함

 

 

2. BranchPythonOperator

def select_random():
	import random
    
    item_lst = ['A', 'B', 'C']
    selected_item = random.choice(item_lst)
    if selected_item == 'A':
    	return 'task_a' # task 1개만 실행시키고 싶은 경우
    elif selected_item in ['B', 'C']:
    	return ['task_b', 'task_c'] # task 2개 이상 실행시키고 싶은 경우
        
python_branch_task = BranchPythonOperator(
	task_id = 'python_branch_task',
    python_callable=select_random
)

python_branch_task >> [task_a, task_b, task_c]

 

 

 

2. BranchPythonOperator 사용 예제코드

'''
dags_branch_python_operator.py
'''

from airflow import DAG
from datetime import datetime
from airflow.operators.python import PythonOperator
from airflow.operators.python import BranchPythonOperator


with DAG(
    dag_id='dags_branch_python_operator',
    start_date=datetime(2023,2,1),
    schedule=None,
    catchup=False
) as dag:
    def select_random():
        import random
        
        item_lst = ['A', 'B', 'C']
        selected_item = random.choice(item_lst)
        if selected_item == 'A':
            return 'task_a' # task 1개만 실행시키고 싶은 경우
        elif selected_item in ['B', 'C']:
            return ['task_b', 'task_c'] # task 2개 이상 실행시키고 싶은 경우
            
    python_branch_task = BranchPythonOperator(
        task_id = 'python_branch_task',
        python_callable=select_random
    )

    def common_func(**kwargs):
        print(kwargs['selected'])

    task_a = PythonOperator(
        task_id = 'task_a',
        python_callable = common_func,
        op_kwargs = {'selected':'A'}
    )

    task_b = PythonOperator(
        task_id = 'task_b',
        python_callable = common_func,
        op_kwargs = {'selected':'B'}
    )

    task_c = PythonOperator(
        task_id = 'task_c',
        python_callable = common_func,
        op_kwargs = {'selected':'C'}
    )

    python_branch_task >> [task_a, task_b, task_c]

 

실행결과 그래프

 

 

실행결과 Xcom