Airflowの実行タスクをBranchPythonOperatorによって分岐させ、結合する方法

# はじめに
Airflowでは、いくつかの方法でタスクの分岐をさせることが可能です。
何かしたの値を受けて、次のいずれかの処理のみを実行させたいときなどに便利ですよね！
今回は`BranchPythonOperator`を使用しようしたタスク分岐の方法と、分岐したタスクを再度結合し、その後の処理を行う方法についてまとめていきます。

# 実行環境
- AWS MWAA環境 (Airflowバージョン2.2.2)

# やってみる
今回は以下の手順で進めていきます。

1. BranchPythonOperatorの基本的な使い方
1. タスクの結合

## BranchPythonOperatorの基本的な使い方
まずは`BranchPythonOperator`の基本的な使い方です。

https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/operators/python/index.html?highlight=branch#airflow.operators.python.BranchPythonOperator

以下のように`python_callable`に分岐させるための関数を渡し、関数側で次のタスクのtask_idを返してやることで簡単に分岐させることができます。
以下の例では、`branch`タスクから`next_true_task`と`next_false_task`に分岐させています。`branch_function`の結果が常に`next_true_task`を返すので、`next_true_task`のみが実行されるはずです。

```python
from airflow import DAG
from airflow.operators.dummy import DummyOperator
from airflow.operators.python_operator import BranchPythonOperator

def branch_function(**kwargs):
    if True:
        return "next_true_task"
    else:
        return "next_false_task"

with DAG(
    "branch_test",
) as dag:
    start = DummyOperator(
        task_id="start",
        dag=dag,
    )

    branch = BranchPythonOperator(
        task_id="branch",
        python_callable=branch_function,
        dag=dag,
    )

    t1 = DummyOperator(task_id="next_true_task")
    t2 = DummyOperator(task_id="next_false_task")

    start  >> branch >> [t1, t2]
```

AirflowUIのGraph Viewを見ると以下のように`branch`タスクから`next_true_task`と`next_false_task`に分岐しているのが分かります。

![スクリーンショット 2022-04-03 16.38.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/643279/2dcfe1a1-623b-3e5d-d40c-e72e78a8cac5.png)


また実際に実行すると、以下のように`next_false_task`が`skipped`(ピンク色)のステータスになりました。問題なさそうです。

![スクリーンショット 2022-04-03 16.42.02.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/643279/3edc610e-aa8a-d215-3574-e9294f464a81.png)


## タスクの結合
タスクの分岐は`BranchPythonOperator`によって簡単に行えましたが、その後再度タスクを合流させ処理を続けたいことがあります。

Graph Viewで見ると以下のような流れですね。新たに`join`、`finish`タスクを置きました。

![スクリーンショット 2022-04-03 16.49.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/643279/758562c2-97f4-74e3-fa87-95329890965e.png)


`join`、`finish`タスクはどちらも`DummyOperator`によるタスクで、中身は空です。
```python
    #
    #
    #
    join = DummyOperator(
        task_id="join"
    )
    finish = DummyOperator(
        task_id="finish"
    )

    start  >> branch >> [t1, t2] >> join >> finish
```

これで問題なさそうなのですが、このまま実行すると最後まで実行されず、以下のように`join`、`finish`タスクまでスキップされてしまいました。

![スクリーンショット 2022-04-03 16.51.17.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/643279/81fc4f25-6d74-ff6a-1813-94d97bac8389.png)


実は各オペレーターの基底クラスである`BaseOperator`には`trigger_rule`という変数があり、以下のような変数になっています。

https://airflow.apache.org/docs/apache-airflow/stable/_modules/airflow/models/baseoperator.html#BaseOperator

```python
"""
    :param trigger_rule: defines the rule by which dependencies are applied
        for the task to get triggered. Options are:
        ``{ all_success | all_failed | all_done | one_success |
        one_failed | none_failed | none_failed_min_one_success | none_skipped | always}``
        default is ``all_success``. Options can be set as string or
        using the constants defined in the static class
        ``airflow.utils.TriggerRule``
    :type trigger_rule: str
"""
```
上流のタスクに自分自身の実行がどのように依存するかを設定するためのもので、全て成功した場合のみの`all_success`や全て失敗した場合のみの`all_failed`など、いくつかの設定値に変更することができます。
そして、これのデフォルト値は`all_success`になっており、上流の全てのタスクが成功しなければなりません。

trigger_ruleの各詳細についてはドキュメントの以下にまとまっています。

https://airflow.apache.org/docs/apache-airflow/stable/concepts/dags.html#trigger-rules

これまでに実装したDAGでは、`join`タスクの上流の`next_true_task`と`next_false_task`はどちらかのみしか実行されず、両方が成功することはありません。そのために、`join`以降のタスクも実行されずスキップされてしまったのでした。

そこで、`join`タスクを以下のように変更します
```python
    #
    #
    join = DummyOperator(
        task_id="join",
        trigger_rule="none_failed_min_one_success",
    )
    #
    #
```

`none_failed_min_one_success`は上流のタスクがひとつも失敗せず、かつ少なくとも一つは成功した場合にトリガーされます。

それでは、DAGを実行してみましょう。
以下のように、`join`、`finish`までタスクが実行されることが確認できました。

![スクリーンショット 2022-04-03 17.10.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/643279/27f6ec4c-bed6-5e83-1e5d-ba44bfe11394.png)


# まとめ
今回は、Airflowにおけるタスクの分岐と結合のやり方をまとめました。
方法自体は単純ですが、ドキュメントにきちんと目を通していかないとハマってしまうポイントかもしれません。

しかし、タスクの分岐、結合など、フローを柔軟に組めるようになると、できることが増えて楽しいですね！
それでは今回はここまで！