Organizando Dags complexos com Task Groups

Criando Pipelines de Dados com Airflow

Volker Janz

Senior Developer Advocate at Astronomer

O problema da complexidade

Dag complexo

Dags grandes ficam difíceis de navegar na visão Graph
Nomes de tarefas se confundem
Pessoas novas no time não acham o que é delas

@task_group

from airflow.sdk import dag, task, task_group

@task_group(
    group_id="ingest_orders",
    default_args={"retries": 3},
)

def process_orders():
    @task
    def extract_orders():
        return [{"id": 1, "amount": 99.99}]

    @task
    def transform_orders(orders):
        return [{"id": o["id"], "total": o["amount"] * 1.08}
                for o in orders]

    return transform_orders(extract_orders())

group_id define um identificador personalizado; por padrão, é o nome da função
default_args se aplicam a todas as tarefas do grupo, evitando repetir configuração
Task groups aparecem como blocos expansíveis na UI

Como aparece na UI do Airflow

Sem task groups

Gráfico simples mostrando três tarefas em linha: extract_orders, transform_orders, load, todas no mesmo nível

Todas as tarefas no mesmo nível

Com task groups

Gráfico mostrando um bloco recolhido chamado process_orders contendo extract e transform, seguido de uma tarefa load fora do grupo

Blocos recolhíveis na UI

Aninhamento e nomes de exibição personalizados

@task_group(group_display_name="Process All")
def process_all():

    @task_group(group_display_name="Ingest Orders")
    def orders():
        return transform(extract())

    @task_group(group_display_name="Process Returns")
    def returns():
        return transform(extract())

    return {
        "orders": orders(),
        "returns": returns(),
    }

Task groups podem ser aninhados para pipelines mais complexos
group_display_name define um rótulo legível na UI (até com emojis)

O padrão factory

@task_group
def process_source(source_name, source_path):

    @task
    def extract():
        return read_data(source_path)

    @task
    def transform(data):
        return clean_data(data)

    return transform(extract())


# Reuse the same pattern for different sources
orders = process_source("orders", "/data/orders.csv")
returns = process_source("returns", "/data/returns.csv")
events = process_source("events", "/data/events.csv")

Task groups são funções Python decoradas
Você pode chamá-las várias vezes com parâmetros diferentes
Esse factory pattern permite reutilizar a lógica do pipeline

Diretrizes para agrupar

Agrupar por domínio ou responsabilidade

Agrupe por domínio ou responsabilidade, não por tipo de operador
Use group_id para nomes programáticos claros e group_display_name para rótulos na UI
Use default_args para compartilhar configs como retries em todas as tarefas do grupo
Aplique a Lei de Miller: mais de 7 itens de topo pode indicar que você precisa de um task group

Vamos praticar!

Criando Pipelines de Dados com Airflow