Apache Airflow 2 overview - part 1

Versions: Apache Airflow 2.4.2 https://github.com/bartosz25/airflow-playground/tree/master/apache-airflow-2/features-part-1

Apache Airflow 2 introduced a lot of new features. The most visible one is probably a reworked UI but there is more! In this and the next blog post I'll show some of the interesting new Apache Airflow features.

Looking for a better data engineering position and skills?

You have been working as a data engineer but feel stuck? You don't have any new challenges and are still writing the same jobs all over again? You have now different options. You can try to look for a new job, now or later, or learn from the others! "Become a Better Data Engineer" initiative is one of these places where you can find online learning resources where the theory meets the practice. They will help you prepare maybe for the next job, or at least, improve your current skillset without looking for something else.

👉 I'm interested in improving my data engineering skillset

See you there, Bartosz

Task Flow API

The first improvement from my list is the Task Flow API. This new method greatly simplifies the DAG declaration by getting rid of many Airflow DSL-related elements. For example, to run an arbitrary Python function you don't need to create a PythonOperator and the associated function. Instead, you can simply define the function and annotate it with a @ŧask decorator.

Want to see it in action? Below you can find a model of the DAG based on the Task Flow API:

    def create_cluster():
        task_definition = {'jar': 'abc.jar'}
        logging.info(f'Creating cluster with {task_definition}')
        return 'cluster_id'

However, there is a catch! The decorator applies to Python-evaluated operators and you can find the full list of the supported elements in the TaskDecoratorCollection:

class TaskDecoratorCollection:
# ...
    python = staticmethod(python_task)
    virtualenv = staticmethod(virtualenv_task)
    external_python = staticmethod(external_python_task)
    branch = staticmethod(branch_task)
    short_circuit = staticmethod(short_circuit_task)

Fortunately, the usage of the Task Flow API allows both old- and new-style DAGs declaration:

    data_availability_checker = PythonSensor(

    def create_cluster():
        task_definition = {'jar': 'abc.jar'}
        logging.info(f'Creating cluster with {task_definition}')
        return 'cluster_id'

data_availability_checker >> create_cluster()

Another improvement this new API brings is the ease of passing the variables between tasks. Since they're just Python functions, the parameters exchange consists of referencing the functions return values:

    def create_cluster():
        task_definition = {'jar': 'abc.jar'}
        logging.info(f'Creating cluster with {task_definition}')
        return 'cluster_id'

    def notify_running_job(cluster_id: str):
        logging(f'A job is running on the cluster {cluster_id}')
        return 'job_id_a'

cluster_creator = create_cluster()
    data_availability_checker >> cluster_creator >> notify_running_job(cluster_creator)

Under-the-hood these "parameters" from the function signature are...XCom variables! The good news is that their exchange is not limited to the @ŧask-decorated functions:

    def notify_running_job(cluster_id: str):
        logging(f'A job is running on the cluster {cluster_id}')
        return 'job_id_a'

    def check_job_status():
        context = get_current_context()
        task_instance: TaskInstance = context['task_instance']
        job_id = task_instance.xcom_pull(task_ids='notify_running_job')
        logging.info(f'Checking job status for {job_id}')
        return True

The last good news is the extendibility. If you need to write your own decorator, you can create a custom one!

Task groups

Another task-related feature in Apache Airflow 2 are task groups. They're a recommended choice over the deprecated SubDAGs to create collections of closely related tasks that should be grouped and displayed together.

Their declaration is relatively easy. It looks like writing a DAG except that each task group can be a parent for another. You can check my example out where I'm generating 12 task groups with 2 subgroups inside each of them:

    with TaskGroup(group_id='partitions_cleaners') as partitions_cleaners:
        months_to_clean = range(12)
        sub_task_groups = []
        for month_to_clean in months_to_clean:
            with TaskGroup(group_id=f'month_{month_to_clean}', parent_group=partitions_cleaners) as month_cleaners:
                partition_1_cleaner = EmptyOperator(task_id=f'merge_{month_to_clean}_partition_1')
                partition_2_cleaner = EmptyOperator(task_id=f'merge_{month_to_clean}_partition_2')
            sub_task_groups.append([partition_1_cleaner, partition_2_cleaner])


Before moving forward, let's stop for a while and see the reasons for depreciating the SubDAGs and preferring task groups:

Task groups repeat the tasks as an instance of the main DAG and overcome the SubDAG issues.

Dynamic task mapping

The 3rd change impacting the DAGs you can write is dynamic task mapping. Sounds mysterious? Actually, the feature is simpler than the name. It lets you generate dynamic DAGs that might have a different number of a given task instance.

In my example below I'm creating a DAG running the refresh_table task as many times as the month day of the execution date:

    def refresh_table(table_name: str):
        logging.info(f'Refreshing {table_name}')
        return table_name

    def generate_params_from_execution_date_day():
        context = get_current_context()
        execution_date: DateTime = context['execution_date']
        logging.info(f'execution_date.day = {execution_date.day}')
        return map(lambda day: f'table_{day}', range(execution_date.day))

    def refresh_materialized_view(refreshed_tables):
        logging.info(f'>>> Refreshing materialized view for tables {", ".join(refreshed_tables)}')

    tables = refresh_table.expand(table_name=generate_params_from_execution_date_day())

Apache Airflow 2 refers to them as to lazy sequences that gets resolved at runtime, only when asked. There is no eager evaluation involved.

Data dependencies

The last feature covered in this first part is about scheduling. Apache Airflow 2 has a possibility to schedule the DAG after a dataset is available. To do that, the new version supports a Dataset trigger. In the example below, I'm generating text files into a watched Dataset directory:

dataset_path = 'file:///tmp/datasets'
dataset = Dataset(uri=dataset_path)

with DAG(dag_id='dataset_writer', start_date=start_date_for_reader_and_writer, schedule='@daily') as writing_dag:

    def data_generator():
        context = get_current_context()
        execution_date: DateTime = context['execution_date']
        with open(f'{dataset_path}/{execution_date.to_date_string()}.txt', 'w') as file_to_write:


with DAG(dag_id='dataset_reader', start_date=start_date_for_reader_and_writer, schedule=[dataset]) as reading_dag:

    def data_consumer():



This data-aware scheduling adds an extra place to the UI where you can visualize all DAGs triggered by each Dataset:

An important thing to notice, though. I haven't succeeded in using the Dataset-scheduling on top of an externally generated data (by not an Airflow-managed writer). Additionally, there is no way to define complex input patterns, such as regular expressions or file glob patterns. That's why you might still need some data sensors in the DAGs to handle these scenarios.

As you can see, Apache Airflow 2 doesn't only come with the new User Interface. It also has several new features extending the tasks and DAG capabilities. And it's only the beginning. I'll share the rest next week in the second part of the blog post.