AIP-76: Hold Dag run until all upstream partitions arrive

[Lee-W]

Why

Closes: #59294

Why

Asset-partitioned Dags that aggregate many upstream slices into one downstream period (e.g., 60-minute-level events rolling up into one hourly Dag run) had no way to express that requirement — the scheduler would fire the downstream run as soon as any single upstream partition arrived.

This PR implements the rollup building block from AIP-76: a Window type that enumerates the full set of upstream partitions required for a downstream period, a RollupMapper that wires a source mapper to a window, and the scheduler logic to gate Dag runs until every required upstream key is present.

What

Partition mappers/windows

• Add Window ABC and six concrete implementations (HourWindow, DayWindow, WeekWindow, MonthWindow, QuarterWindow, YearWindow) to both airflow-core and the Task SDK
• Add RollupMapper that composes a source_mapper with a Window and exposes to_upstream(downstream_key) → frozenset[str]
• Add decode_downstream / encode_upstream hooks to PartitionMapper and implement them in _BaseTemporalMapper; StartOfWeekMapper gets a regex-based override because %V is ambiguous with strptime.
• Add week_start parameter to StartOfWeekMapper for non-Monday week starts

Scheduler

• Rewrite _create_dagruns_for_partitioned_asset_dags to bulk-fetch serialized Dags and partition-key logs, removing N+1 queries, and cap per-tick work at MAX_PARTITION_DAG_RUNS_PER_TICK
• Add _resolve_asset_partition_status / _check_rollup_asset_status to evaluate rollup satisfaction; non-rollup assets continue to satisfy immediately

Serialization

• Add encode_window / decode_window and extend mapper encoder/decoder to round-trip RollupMapper and all Window subclasses

UI / API

• Enrich next_run_assets endpoint with per-asset received_count, required_count, received_keys, required_keys, and is_rollup for partitioned Dags
• Update AssetNode and AssetSchedule components to surface rollup progress (e.g. "12 / 24 received")
• Add AssetProgressCell for inline progress in the Dags list

---

Was generative AI tooling used to co-author this PR?

• Yes — Claude Sonnet 4.6

Generated-by: Claude Sonnet 4.6 following the guidelines

with DAG(
    dag_id="daily_team_a_rollup",
    schedule=PartitionedAssetTimetable(
        assets=team_a_player_stats,
        default_partition_mapper=RollupMapper(
            source_mapper=StartOfDayMapper(),
            window=DayWindow(),
        ),
    ),
    catchup=False,
    tags=["player-stats", "rollup"],
):
    """
    First rollup level: 24 hourly partitions of ``team_a_player_stats`` → one daily summary.

    ``StartOfDayMapper`` normalizes each upstream hourly timestamp (``%Y-%m-%dT%H:%M:%S``)
    to its day-start (``%Y-%m-%d``); ``DayWindow`` declares the downstream run needs
    all 24 hourly partitions before firing. Publishes ``daily_team_a`` so the
    monthly rollup below can consume it.
    """

    @task(outlets=[daily_team_a])
    def summarise_team_a_day(dag_run=None):
        """Produce the full-day rollup once every hour has arrived."""
        if TYPE_CHECKING:
            assert dag_run
        print(f"All 24 hourly partitions received. Day: {dag_run.partition_key}")

    summarise_team_a_day()

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[Copilot]

Pull request overview

Implements “rollup” support for partition mappers (AIP-76) by introducing a RollupMapper interface with to_upstream() and using it in the scheduler to wait for a complete set of upstream partition keys before creating partitioned asset-triggered DAG runs.

Changes:

• Add RollupMapper base class (core + task SDK) with an abstract to_upstream() contract.
• Implement to_upstream() for weekly and monthly temporal mappers (core + task SDK).
• Update the scheduler’s partitioned-asset DAG-run creation logic to enforce rollup completeness when a mapper supports it.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py	Introduces SDK-side RollupMapper abstraction.
task-sdk/src/airflow/sdk/definitions/partition_mappers/temporal.py	Adds SDK to_upstream() for week/month temporal rollups.
airflow-core/src/airflow/timetables/base.py	Adds get_partition_mapper() hook to the Timetable protocol.
airflow-core/src/airflow/partition_mappers/base.py	Introduces core-side RollupMapper abstraction.
airflow-core/src/airflow/partition_mappers/temporal.py	Adds core to_upstream() for week/month temporal rollups.
airflow-core/src/airflow/jobs/scheduler_job_runner.py	Uses rollup mapper behavior to decide when partitioned asset-triggered DAG runs are ready.

[Lee-W]

from __future__ import annotations

from airflow.sdk import (
    DAG,
    Asset,
    CronPartitionTimetable,
    PartitionedAssetTimetable,
    WeeklyRollupMapper,
    task,
)

daily_sales = Asset(uri="file://incoming/sales/daily.csv", name="daily_sales")

# Upstream Dag: produces one partition per day (key format: "2024-01-15T00:00:00")
with DAG(
    dag_id="ingest_daily_sales",
    schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"),
):

    @task(outlets=[daily_sales])
    def ingest():
        pass

    ingest()


# Downstream Dag: runs once all 7 daily partitions for a week have arrived
with DAG(
    dag_id="weekly_sales_report",
    schedule=PartitionedAssetTimetable(
        assets=daily_sales,
        default_partition_mapper=WeeklyRollupMapper(),
    ),
    catchup=False,
):

    @task
    def generate_report(dag_run=None):
        # dag_run.partition_key will be the week key, e.g. "2024-01-15 (W03)"
        print(dag_run.partition_key)

    generate_report()

[Lee-W]

The backend part is basically wrapped up, but the frontend and API side need some work. The UI is quite weired for these cases now

[Copilot]

Pull request overview

Copilot reviewed 21 out of 21 changed files in this pull request and generated 13 comments.

Comments suppressed due to low confidence (1)

airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py:1

• Counting PartitionedAssetKeyLog.id can over-count when duplicate log rows exist for the same upstream partition key (e.g. retries/dup inserts), inflating total_received and potentially showing the run as satisfiable earlier than it should be. Consider counting distinct PartitionedAssetKeyLog.source_partition_key (and/or a distinct composite of (asset_id, source_partition_key)) to match the scheduler’s set-based satisfaction semantics.

# Licensed to the Apache Software Foundation (ASF) under one

[uranusjr]

I do’t like we need to load the timetable here… Is there a way to pre-calculate this somewhere?

[Lee-W]

what do you think if we save this to DagModel.partition_mapper_info?

 {
    "name, uri": {
      "is_rollup": true,
      "runtime_partitioned": false,
    },
    "name, uri": {
      "is_rollup": false,
      "runtime_partitioned": false,
    }
  }

something like this?

[uranusjr]

Same problem

[uranusjr]

What’s the worst can happen if this is not the best value? (In other words, should we provide a way to change this?)

[Lee-W]

What’s the worst can happen if this is not the best value

The DagRun that waits for 501 Assets will never be triggered.

should we provide a way to change this?

I think we should, but I'm thinking of doing it in another PR.