Airflow Fundamentals study materials
Airflow Fundamentals study resources

A PythonOperator That Works Alone and Fails in the DAG

Airflow Fundamentals · Put your knowledge to work in context.

All case studies

A PythonOperator That Works Alone and Fails in the DAG

An engineer writes a three-task DAG: extract returns a dict of row counts, transform is supposed to use those counts to decide how many partitions to build, and load writes the result. Each function is tested locally by calling it directly and all three pass. In Airflow, transform raises NameError: name 'counts' is not defined, and when the engineer 'fixes' it by making counts a module-level global, the value is always None on the second day's run even though extract logged the correct numbers. Separately, extract returns a dict containing a datetime and a Decimal, and the engineer is considering passing the entire dict through XCom as-is. The cluster runs the CeleryExecutor with two workers on different hosts.

1. Explain both failures - the NameError and the stale None - and describe how data should actually flow from `extract` to `transform`.

Read the worked answer

Both failures come from treating a task function as an ordinary in-process function call. In Airflow, each task runs in its own process, potentially on a different worker host, and the only supported channel for passing data between tasks is XCom. transform has no counts parameter and no closure over extract's local variable, hence the NameError.

The module-level global 'fix' fails for a different reason: the global is set in the worker process that ran extract, but transform may run in a different process, a different worker, and on the next day a freshly started one. Nothing carries the value across. The None is whatever the global's initial value was in the process that happened to run transform. This is exactly the class of bug that CeleryExecutor with multiple hosts makes obvious and a single-process LocalExecutor can mask.

The correct pattern is explicit:

  • extract returns the value; Airflow pushes the return value to XCom under key return_value.
  • transform declares a parameter (e.g. def transform(counts, **context)) and the DAG wires transform(counts=extract_task.output) in TaskFlow style, or the function calls context['ti'].xcom_pull(task_ids='extract') in classic style.
  • Keep the payload small: pass counts, keys, or a path, not the data itself. Large XComs bloat the metadata database and are serialized/deserialized on every access.
  • If the real payload is large, write it to S3 or a staging table in extract and pass only the URI through XCom.

2. The engineer passes the dict containing a `datetime` and a `Decimal` through XCom and it fails. Explain why, and give two robust alternatives.

Read the worked answer

XCom values are serialized for storage in the metadata database. The default serializer is JSON-based (Airflow ships a JSON serializer that handles some extra types), and arbitrary Python objects such as datetime and Decimal are not reliably round-trippable. Depending on configuration the push fails at serialization, or the pull returns a string or a dict with unexpected types, and transform then breaks on arithmetic. The failure is configuration-dependent, which is why it can appear to work on one deployment and fail on another.

Robust alternatives:

  1. Pass primitives only. Convert at the boundary: {'rows': int(n), 'as_of': dt.isoformat(), 'amount': str(dec)}, and convert back in the consumer. This is explicit, portable, and survives any serializer change.
  2. Pass a reference, not the data. Have extract write a JSON or Parquet artifact to object storage or a staging table and push only the key/URI. The consumer reads and parses it with full control over types. This also solves the size problem.

A third option - enabling a custom XCom backend (e.g. one backed by S3) with a serializer that handles these types - is legitimate but is an infrastructure change affecting every DAG, so it should be a deliberate platform decision, not a local workaround. The key judgement is that XCom is a small-message channel between tasks, not a general-purpose data bus, and type discipline at the boundary is cheaper than debugging serializer behaviour.

3. `extract` occasionally takes 40 minutes. A colleague suggests wrapping the whole DAG in a single PythonOperator that calls all three functions in sequence, 'to avoid XCom entirely'. Assess this.

Read the worked answer

It removes the XCom problem by removing the task graph, and that is the wrong trade for this pipeline.

Losses:

  • No partial retry. If load fails after extract and transform succeeded, the retry re-runs the 40-minute extract. With three tasks, retries on load alone re-runs only the cheap step.
  • No visibility. The UI shows one green or red box. You cannot see which stage is slow or which failed without reading logs.
  • No parallelism or fan-out. A single task cannot branch, cannot run partitions concurrently via dynamic task mapping, and cannot express dependencies on other DAGs per stage.
  • No independent scheduling or SLAs per stage.
  • Worker slot held for the entire duration with no checkpointing between stages.

Legitimate uses of a single operator exist - a genuinely atomic step, or a thin wrapper around a script - but a three-stage ELT with a 40-minute extract is not one.

If the concern is XCom ergonomics, the better answer is TaskFlow: transform(extract()) reads like a function call while still producing three real tasks with real retries and real visibility. If the concern is XCom size, pass a URI. The colleague's suggestion solves a small serialization annoyance by discarding retry granularity and observability, which are the main reasons to use Airflow at all.