2
votes

I figure out that xcom actually write data into database and pull it from other task. My dataset is large and pickle it and write to database cause some unnecessary delay. Is there a way to communicate data between tasks in the same airflow Dag without using xcom?

below is the code that I tried, the context is actually not passed. I know use task_instance.xcom_push() would work but it also pickle the data and write it to database which I don't need.

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime, timedelta
from pandas import DataFrame
import pandas as pd
from custom.dataframe_to_postgres_operator import PostgresOperatorBulk
from airflow.operators.postgres_operator import PostgresOperator

def read_df(task_instance, **context):
    df = pd.read_parquet('/usr/local/airflow/data/df.parquet.gzip')
    print(df)
    # task_instance.xcom_push('data', df)
    context.update({'data': df})
    for k, v in context.items():
        print(k, v)
    return 1

def get_df(task_instance, **context):
    for k, v in context.items():
        print(k, v)
    df = context['data']

default_args = {
    'owner': 'Airflow',
    'depends_on_past': False,
    'start_date': datetime(2020, 2, 17),
    'retries': 0,
}

dag = DAG('abcdefg', default_args=default_args, schedule_interval=timedelta(days=1))

task_read_df = PythonOperator(
    task_id='read_df',
    python_callable=read_df,
    dag=dag,
    provide_context=True,
    do_xcom_push=False
)

task_get_df = PythonOperator(
    task_id='get_df',
    python_callable=get_df,
    dag=dag,
    provide_context=True,
    do_xcom_push=False
)

task_read_df >> task_get_df
1

1 Answers

1
votes

If you have a large dataset that you want to exchange I suggest storing the data in some form of a temporary location (e.g. a designated directory) and then passing the path to such a temp file or files using XCOM (which for small data pieces is cheap and gives good enough performance).

For that, a good library is tempfile which helps ease the pain to avoid duplicates among the temp files.

Why XCOM and not shared execution context between Tasks

Given that the Tasks can be executed in parallel this is the first issue that creates a lot of difficulty for Python (GIL, sharing data in parallel).

Secondly, to ensure some form of persistence (and therefore resilience to failure) you have to use a database to ensure ACID.

It all makes XCOM mechanism relatively heavy (particularly when you add pickling on top of it) but it is universal.

Having all that in mind you have to remember that usage of temporary files to which you pass a path via XCOM does ensure the same level of resilience as XCOM itself (particularly if the file is stored on RAM disk). It also does not support replay of Tasks unless you keep the temp files indefinitely.