Best practices for implementing data-aware DAGs using Datasets or Dataset aliases

[pedro-cf]

I'm looking to implement a data pipeline using two DAGs and Airflow's Datasets feature. I'd appreciate guidance on the best way to structure this using Datasets or Dataset aliases.

Current setup:

1. ingest_DAG:

   • Fetches data from an external API
   • Processes and loads the data into a PostgreSQL table (ingest_table)
   • Includes a run_id column in ingest_table to uniquely identify each batch of ingested data

2. transform_DAG:

   • Triggered by the Dataset produced by ingest_DAG
   • Must receive the run_id to identify the specific batch of data to process
   • Reads data from ingest_table using the received run_id
   • Processes the data
   • Inserts results into transform_table

Current solution:

# ingest_dag.py
from airflow import DAG
from airflow.decorators import task
from airflow.datasets import Dataset, DatasetAlias
from datetime import datetime, timedelta

DATA_ALIAS = DatasetAlias("data_alias")

with DAG(
    'ingest_dag',
    start_date=datetime(2023, 1, 1),
    schedule="@hourly",
    catchup=False,
    is_paused_upon_creation=False,
) as dag:

    @task
    def fetch_data():
        return {"data": "sample data"}

    @task
    def load_to_database(data):
        print(f"Loading data to database: {data}")

    @task(outlets=[DATA_ALIAS])
    def produce_dataset(data, run_id=None, outlet_events=None):
        print(f"Producing dataset for run_id: {run_id}")
        outlet_events[DATA_ALIAS].add(Dataset("data_alias"), extra={"run_id": run_id})

    fetched_data = fetch_data()
    loaded_data = load_to_database(fetched_data)
    produce_dataset(loaded_data)

# transform_dag.py
from airflow import DAG
from airflow.decorators import task
from airflow.datasets import Dataset, DatasetAlias
from datetime import datetime, timedelta

DATA_ALIAS = DatasetAlias("data_alias")

with DAG(
    'transform_dag',
    start_date=datetime(2023, 1, 1),
    schedule=DATA_ALIAS,
    catchup=False,
    is_paused_upon_creation=False,
) as dag:

    @task(inlets=[DATA_ALIAS])
    def process_data(inlet_events=None):
        events = inlet_events[DATA_ALIAS]
        triggering_run_id = events[-1].extra["run_id"]
        print(f"Processing data for run_id: {triggering_run_id}")
        return {"processed_data": "sample processed data"}

    @task
    def store_processed_data(data):
        print(f"Storing processed data: {data}")

    processed_data = process_data()
    store_processed_data(processed_data)

Questions:

1. Is my current approach using Datasets and Dataset aliases effective for this pipeline, or are there improvements I could make?
2. In my current setup, am I correctly passing the run_id from ingest_DAG to transform_DAG using Datasets?
3. I'm currently including the run_id in the Dataset's extra field. Is this the recommended way to pass metadata between DAGs, or are there better alternatives?
4. In the transform_DAG, how can I ensure that only the data corresponding to the received run_id is processed?
5. Given my current implementation, are there any potential pitfalls or common mistakes I should be aware of when using Datasets for inter-DAG communication?
6. Are there any performance considerations I should be aware of with this approach, especially as the volume of data increases?
7. How can I ensure idempotency in my transform_DAG, particularly if it needs to be rerun for the same run_id?
8. Are there any Airflow features or patterns that could simplify or improve my current implementation?

Any insights, examples, or recommendations based on my current solution would be greatly appreciated. Thank you!

[pedro-cf]

Greetings @jedcunningham any chance you have an update for this answer since #40856 (reply in thread) :)