AIP-76 Asset Partitions

Status

This AIP is part of AIP-73, which aims to expand data awareness within Airflow.

Motivation

Partitions are a critical enhancement to Airflow's data processing capabilities, designed to address some limitations in the current task-centric approach. The primary motivations for introducing partitions are as follows:

• Granular Data Management: Currently, Airflow's data intervals are tied directly to task schedules, which limits the flexibility in handling datasets that do not fit neatly into these intervals. By introducing partitions, we can manage data at a more granular level, allowing for efficient processing of specific slices of data. This is particularly useful for large datasets that need to be processed incrementally or updated periodically.
• Enhanced Performance: Processing large datasets in their entirety can be resource-intensive and time-consuming. Partitions enable Airflow to handle smaller chunks of data independently, improving performance and reducing processing times. This is essential for maintaining the efficiency and scalability of data flows, especially as data volumes continue to grow.
• Improved Data Lineage and Impact Analysis: With partitions, it becomes easier to track the flow of data through various stages of the workflow. This visibility allows for more accurate impact analysis when changes are made, ensuring that any modifications to upstream data can be traced and addressed effectively. This capability is crucial for maintaining data quality and reliability. 
• Facilitating Incremental Processing: Many modern data processing strategies rely on incremental updates, where only new or changed data is processed. Partitions support this approach by allowing for the independent processing of data slices. This facilitates more efficient and intelligent backfilling, enabling data flows to quickly catch up on missed data without reprocessing the entire dataset.
• Flexibility in Data Organization: Different user-cases may require different methods of data partitioning, such as time-based or segment-based partitions. By providing a flexible partitioning system, Airflow can better accommodate a variety of data organization strategies, supporting more complex and nuanced data flows.

Rationale

The idea behind partitions is to provide a way to model and organize data that is closely related, but still distinct, decoupling it from the task. For example, a daily-partitioned table for sales orders could be seen as a collection of partitions that correspond to a sales order data asset corresponding to sales orders made on each day. These partitions are computed using the same code and derived from the same upstream sources. They can be processed and visualized in bulk and individually or as subsets.

Partitions in Airflow are a logical concept and don’t necessarily need to correspond to storage-level partitions. Some systems like Hive have an explicit notion of partitions with table partitions corresponding to subdirectories, while other systems like Snowflake abstract storage away entirely. Even without a direct link, the concept is incredibly useful for data orchestration. For instance, when overwriting data in a Snowflake table for a specific date range, you delete all the rows for that date and insert a new set of rows — essentially working with a logical partition.

Consider there’s an asset populated by a function. If the asset is not partitioned, the function can be thought to logically replace the entire asset every time it writes.

Although history is always kept, the previous execution result is logically overwritten when a new run happens.

If the asset is partitioned, on the other hand, every execution of the function would receive a key to identify a section of the asset it should update. An asset partitioned by date, for example, can have the function receive the date as the partition key. Execution then looks like this:

In this simple form, the partition key is actually just the logical date we already have in Airflow. However, partitions are fundamentally a property of the data, not the task or DAG. Airflow’s current implementation has constraints:

1. The data interval or partition is automatically calculated based on the schedule and requires all data updated by different tasks within the DAG to be partitioned in the same way.
2. It requires for partitions to be one-to-one with task runs.
3. The partition key must be a date on a pre-defined schedule. This is fine for many ETL use cases, but problematic in edge or non-ETL cases.

A partition describes how an asset is partially generated whenever its function executes. While the function itself does not strictly require this information (the asset’s function body can do whatever it wants and is the canonical source of truth for what gets written), there are clear benefits to using partitions:

1. Clarity in Asset Generation: Partitions provide other parts of the system with information on how an asset is generated without analyzing data outputs. For example,  if the raw sales order asset uses hourly partitions and the daily sales order asset depends on the raw sales order asset, we know that 24 hourly partitions from the raw sales order asset correspond to 1 daily partition in the daily sales order asset. This clarity allows for a more granular highlight of data flow in the UI and showcases how issues impact downstream assets.
2. Parallel Processing: Multiple partitions can be processed in parallel. For example, if partitions represent cloud regions, we can launch individual DAG runs to process data for each cloud region more quickly. Additionally, if upstream data has been processed for some regions but not others, processing can start for the regions that are ready.
3. Ad-hoc Range Processing: A single task can generate a range of partitions ad-hoc, such as during backfills. If an asset is scheduled to run once a day and normally processes a day's worth of data, it may be more efficient to delegate the parallelization to a compute engine like Spark and kick off a single task covering multiple partitions. Using partitions allows for this capability, enhancing performance and flexibility.

Specification

Partitions can be defined statically or dynamically. Static partitions can also be combined to create more complex partitioning. Partition information can be accessed inside asset functions that.

A sliced partition of an asset is identified by a value of any Airflow-serializable types. The partition’s identity (within the same asset) relies on the value’s hash() result and implementation of __eq__.

Two categories of partitions will be introduced: time-based, and segment-based. A time-based partition is specified an interval, e.g. daily, hourly, a cron expression, or timedelta. A segment-based partition is specified by a static list.

Time-based Partitions

These are similar to logical dates (more accurately, data intervals) today, and are meant to represent time or date ranges. Like today, they can automatically be derived from the schedule or can be manually set.

@asset(..., schedule="@hourly", partition=PartitionByInternal("@hourly"))
def hourly_data():
    ...

The concept is similar to how Airflow 2 recommends using the data interval, but they are fundamentally different since the partition boundaries are entirely decoupled from the schedule. Consider this example:

@asset(..., schedule="@daily", partition=PartitionByInternal("@hourly"))
def hourly_data_generated_daily():
    ...

The asset is materialized once per day at midnight, but every materialisation would trigger 24 runs, each covering one hour of the previous day.

Conversely, if an hourly-scheduled asset has a daily partition, it would materialize once per hour, but for any given day, the first 23 materialization runs end immediately without actually executing the function (the task being skipped), only the last one being actually executed to cover the past entire day. This is useful if the asset follows another asset’s schedule, but does not want to be materialized as often. See Partition Dependency below.

Segment-based Partitions

These are meant to represent other types of logical groupings, like teams or cloud regions. For example, if an asset is generated every day to evaluate the spend of multiple data warehouses, each data warehouse can be represented as a partition.

@asset( 
    ..., 
    schedule="@daily",
    partition=PartitionBySequence(["marketing-dwh", "engineering-dwh"]),
)
def dwh_daily_cloud_spend():
    ...

Combined Partitions

You will also be able to combine time-based and segment-based partitions.

@asset(
    schedule=dwh_daily_cloud_spend, 
    partition=PartitionByProduct(
        [PartitionByInternal("@hourly"), PartitionBySequence(["marketing-dwh", "engineering-dwh"])],
    ),
)
def dwh_spend_analysis():
   ...

After the upstream asset is generated each day, 48 (24 hours x 2 data warehouses) runs are triggered. To simplify logic, a combined partition definition can only contain at most one time-based partition definition. We’ll explore more complex combinations if there are concrete real-world use cases after this feature is fully implemented and rolled out.

Runtime-assigned Partitions

All cases discussed up until now know the partition key up-front, and therefore can provide the key to a particular function run. Sometimes, however, partitions are not known until when the function is actually run to process the incoming data. For example, to predict energy consumption, an energy company trains an ML model per customer. Customers are frequently added and removed, so it is simpler to get an accurate list of customers at runtime.

active_customers = PartitionAtRuntime(name="active_customers")

@asset(schedule="@hourly", partition=active_customers)
def energy_consumption_model():
    """get customer list and train model"""
    customers = ... # get customer list
    for customer in customers:
        train_model = ...
        active_customers.add_partition(key=customer)

Runtime partitioning cannot be combined with statically defined partitioning with PartitionByProduct.

Rich Partition Key

Internally, each partition is identified by its partition key. Although a primitive value (like a date, or a string) can satisfy basic use cases, real-world scenarios most likely require richer information, such as a time span (e.g. 2pm yesterday to 2pm today) or a collection of values (e.g. x, y, and z in a three-dimension coordinate system). We also need to consider how we can disply these in the web UI. The partition key, either generated automatically from a definition or given directly by the user (with either PartitionBySequence or added at runtime), with always be wrapped in a PartitionKey class. The user can create custom subclasses to contain custom fields, like this:

class Position(PartitionKey):
    x: float
    y: float
    z: float

    # Optional; Airflow generates a repr similar to dataclasses.
    def __str__(self) -> str:
        return f"({self.x}, {self.y}, {self.z})"

The PartitionKey type should have a custom metaclass to stop users from being able to supply a custom __init__ or __new__ to inject arbitrary code into protected Airflow processes.

In simple cases (e.g. user supplying plain strings to PartitionByProduct), Airflow should automatically wrap the keys in special classes, and unwrap them automatically when possible. PartitionByInternal should produce Interval instances as partition keys that specify the start and end of the partition, with a custom __str__ that represents the interval by the start date.

Accessing Partition Information

Similar to logical dates and data intervals, partitions are parameterizable and will be available via the task context.

@asset(
   ...,
   partition=PartitionBySequence(["marketing-dwh", "engineering-dwh"]),
)
def dwh_daily_cloud_spend():
    context = get_current_context()
    for partition in context["partition"].iter_past_partitions():
        # do something

The partition value in context is a rich object that exposes various partition information, including the original definition, the current and past partitions, etc. This has not been fully hashed out yet.

Partition Dependency

When an asset depends on an upstream asset, and they have the same partitioning logic, each partition of the downstream depends on the corresponding partition of the upstream.

@asset(schedule="@hourly", partition=PartitionByInternal("@hourly"))
def raw_hourly_data():
    ...

@asset(schedule=raw_hourly_data, partition=PartitionByInternal("@hourly"))
def hourly_data():
    ...

The partition of raw_hourly_data covering 2pm must be materialized first for the partition of hourly_data covering the same period to be materialized.

If the assets’ partitions do not align exactly, all upstream partitions that overlaps the downstream must be materialized first. So

@asset(schedule=raw_hourly_data, partition=PartitionByInternal("30 * * * *"))
def hourly_data_shifted():
    ...

To materialize partition covering 14:30–15:30, both the 2pm and 3pm partitions of raw_hourly_data must first materialize.

If you want a downstream to aggregate multiple partitions from the upstream, you can do

@asset(schedule=hourly_data, partition=PartitionByInternal("@daily"))
def aggregated_daily_data():
    ...

Every partition of this asset depends on 24 partitions of hourly_data of the day.

Conversely, if an upstream is not updated often, but the downstream needs to be calculated more often, you can do

@asset(schedule="@yearly", partition=PartitionByInterval("@yearly"))
def yearly_data():
    ...

@asset(schedule=yearly_data, partition=PartitionByInternal("@monthly"))
def data_using_slow_data():
    ...

After each yearly update, 12 materialization runs are scheduled for the downstream since they all depend on the same upstream partition.

Backfilling

With the context of partitions, backfilling is an operation that retroactively materializes time-based partitions in the past. Only assets with time-based partitioning can be backfilled—the concept does not make sense otherwise (assets either partitioned by sequence or not at all).

An asset also needs to provide the earliest partition to backfill from. This can be done by the start_date argument on @asset. This argument is optional—if not given, the asset cannot be backfilled, and materializations are scheduled only for later dates. We’re not adding a separate catchup argument to assets. The behavior is always the same as catchup=False; backfills must be initiated by a user request after an asset is created.

This AIP does not intend to design how backfills are implemented (should run via the executor and things like that). Here we only intend to understand what consists of a backfill in an abstract sense.

Manual Runs

Asset materialization can be triggered manually. Each manual trigger needs to specify exactly which partition to cover. We should have a rich selector mechanism in the web UI. Re-materializing an existing partition should be considered a re-run and subject to AIPs 63, 64 and 66—an asset is internally analogous to a DAG, and its materialization recorded with TaskInstance. If a partition has previously been successfully materialized manually, later scheduled materialization will skip that partition.

Future Work

Topics listed in this section are incremental enhancements to the core concepts, and can be implemented at any point without breaking compatibility. They may still be implemented and released with other parts of this AIP if we can figure out the needed details in the process, but are not requirements for this AIP to be considered DONE.

Shorthands for Time-based Partitioning

There are some possibilities to simplify the partition argument. For example, "@daily" seems like an obvious shorthand to PartitionByInternal("@daily"). This does not work as well with more sophisticated cases, such as

@asset(..., partition=PartitionByInternal(
    "@daily",
    start=datetime.date(2024, 1, 1),  # For backfilling.
))
def daily_data_from_2024():
    ...

It is not possible to express this with simple Python primitives. Since arguably most real-world scenarios require backfilling to some extent, the shorthand suddenly not too worthwhile if it only works for simplified tutorials.

There are also ambiguous cases. Shortening PartitionBySequence(["foo", "bar"]) to ["foo", "bar"] looks obvious at first glance, but it also makes sense to shorten a PartitionByProduct to a list. Since it is technically not possible to reliably distinguish the two cases, maybe the best approach to do neither in the name of explicitness.

Automatically Partition Inference

A lot of the time, we just want to partition the asset by the schedule—for example, if the asset is materialized every day at noon, a partion covers the previous noon to the lastest. It would be useful to not need to repeat the 0 12 * * * twice. Similarly, if an asset does not have a schedule (only materialized manually), it normally isn’t partitioned either. It would be nice if we can reduce the boilerplate.

One idea is to make partition optional, and defaults to “auto mode” that:

• If the schedule is time-based, a partition is created from the period between each scheduled run.
• If the schedule is not time-based, there is no partition.

This sounds reasonable, but how about assets that depend on a time-based asset? Do they automatically use the same partitions? What if they depend on multiple assets with different schedules?

A quick experiement also revealed a minor reability issue. If I see

@asset(
    schedule="@daily",
    # Many
    # lines
    # of
    # params
    # ...
)
def my_asset():
    ...

I must read all the lines carefully to find if there’s an explicit partition parameter, and can be sure the partition is @daily only after reading every simple line and making sure partition is not in there. If the argument is required, I can always simply scan for partition= to find the information I need.

We are therefore making partition a required parameter in the initial proposal. This buys us time to know what is the best solution here by observing real use cases.

Complex Schedule-Partition Mapping

Up until now, we’ve only considered time-based cases where the partition logically related to the materialization time. Sometime, however, you would want even more interesting partitions. For example, a hourly-materialized asset that modifies data for the past two hours (the new hour, and the hour before it that was also covered by the previous partition—a.k.a. a slicing window). There are even more exotic use cases, and it is not possible for Airflow to include every possibility. It is likely better to have a pluggable interface similar to Timetable for users to provide custom partitioning logic, and Airflow can adopt popular ones.

Complex Partition Dependency Rules

Similar to the schedule-partition mapping, we also only have considered partition dependencies where the upstream partitions cover downstream ones. There are, however, valid uses where you want to depend on other upstream partitions. Again, it is better to allow for custom dependency rules. The situation here is somewhat alike trigger rules—the ability to customise it is also raised as a feature request from time to time. They are also similarly non-trivial and requiring more thought.