# 2.4 Branch Operator

The `BranchOperator` is used to decide which path to run based on the input data. 
Example, if you have two paths, you can decide which path to run based on the input data.

There are two ways to use the `BranchOperator`:

## Build A `BranchOperator` With A Branch Mapping

Pass a dictionary of branch functions and task names to the `BranchOperator` constructor.

```python
from dbgpt.core.awel import DAG, BranchOperator, MapOperator

def branch_even(x: int) -> bool:
    return x % 2 == 0

def branch_odd(x: int) -> bool:
    return not branch_even(x)

branch_mapping = {
    branch_even: "even_task",
    branch_odd: "odd_task"
}

with DAG("awel_branch_operator") as dag:
    task = BranchOperator(branches=branch_mapping)
    even_task = MapOperator(
        task_name="even_task", 
        map_function=lambda x: print(f"{x} is even")
    )
    odd_task = MapOperator(
        task_name="odd_task", 
        map_function=lambda x: print(f"{x} is odd")
    )
```

In above example, the `BranchOperator` has two child tasks, `even_task` and `odd_task`. 
The `BranchOperator` will decide which child task to run based on the input data.
So we pass a dictionary of branch functions and task names to the `BranchOperator` 
constructor to define the branch mapping, in dictionary, the key is the branch function,
and the value is the task name, when run the branch task, all the branch function will 
be executed, if the branch function return `True`, the task will be executed, otherwise,
 it will be skipped.

## Implement A Custom `BranchOperator`

Just override the `branches` method to return a dictionary of branch functions and task names.

```python
from dbgpt.core.awel import DAG, BranchOperator, MapOperator

def branch_even(x: int) -> bool:
    return x % 2 == 0

def branch_odd(x: int) -> bool:
    return not branch_even(x)

class MyBranchOperator(BranchOperator[int]):
    def __init__(self, even_task_name: str, odd_task_name: str, **kwargs):
        self.even_task_name = even_task_name
        self.odd_task_name = odd_task_name
        super().__init__(**kwargs)
        
    async def branches(self):
        return {
            branch_even: self.even_task_name,
            branch_odd: self.odd_task_name
        }

with DAG("awel_branch_operator") as dag:
    task = MyBranchOperator(even_task_name="even_task", odd_task_name="odd_task")
    even_task = MapOperator(
        task_name="even_task", 
        map_function=lambda x: print(f"{x} is even")
    )
    odd_task = MapOperator(
        task_name="odd_task", 
        map_function=lambda x: print(f"{x} is odd")
    )
```

## Examples

### Even Or Odd

Create a new file named `branch_operator_even_or_odd.py` in the `awel_tutorial` directory and add the following code:

```python
import asyncio
from dbgpt.core.awel import (
    DAG, BranchOperator, MapOperator, JoinOperator, 
    InputOperator, SimpleCallDataInputSource,
    is_empty_data
)

def branch_even(x: int) -> bool:
    return x % 2 == 0

def branch_odd(x: int) -> bool:
    return not branch_even(x)

branch_mapping = {
    branch_even: "even_task",
    branch_odd: "odd_task"
}

def even_func(x: int) -> int:
    print(f"Branch even, {x} is even, multiply by 10")
    return x * 10

def odd_func(x: int) -> int:
    print(f"Branch odd, {x} is odd, multiply by itself")
    return x * x

def combine_function(x: int, y: int) -> int:
    print(f"Received {x} and {y}")
    # Return the first non-empty data
    return x if not is_empty_data(x) else y
    
with DAG("awel_branch_operator") as dag:
    input_task = InputOperator(input_source=SimpleCallDataInputSource())
    task = BranchOperator(branches=branch_mapping)
    even_task = MapOperator(task_name="even_task", map_function=even_func)
    odd_task = MapOperator(task_name="odd_task", map_function=odd_func)
    join_task = JoinOperator(combine_function=combine_function, can_skip_in_branch=False)
    input_task >> task >> even_task >> join_task
    input_task >> task >> odd_task >> join_task

print("First call, input is 5")
assert asyncio.run(join_task.call(call_data=5)) == 25
print("=" * 80)
print("Second call, input is 6")
assert asyncio.run(join_task.call(call_data=6)) == 60
```
Note: `can_skip_in_branch` is used to control whether current task can be skipped in the branch. 
Set it to `False` to prevent the task from being skipped.

And run the following command to execute the code:

```bash
poetry run python awel_tutorial/branch_operator_even_or_odd.py
```

And you will see the following output printed to the console.

```bash
First call, input is 5
Branch odd, 5 is odd, multiply by itself
Received EmptyData(SKIP_DATA) and 25
================================================================================
Second call, input is 6
Branch even, 6 is even, multiply by 10
Received 60 and EmptyData(SKIP_DATA)
```

The graph of the DAG is like this:


<p align="left">
  <img src={'/img/awel/awel_tutorial/branch_operator_example_1.png'} width="1000px"/>
</p>


In above example, the `BranchOperator` has two child tasks, `even_task` and `odd_task`, 
it will decide which child task to run based on the input data and the branches mapping.

We also use the `JoinOperator` to combine the data from both child tasks, if a path is 
skipped, the `JoinOperator` will receive an `EmptyData(SKIP_DATA)` as input data, and we
can use `dbgpt.core.awel.is_empty_data` to check if the data is empty data.
