Expose upstream partition key on triggering asset events

[Lee-W]

Why

• A downstream consumer task couldn't tell which upstream partitions fed its run — triggering_asset_events (and dag_run.consumed_asset_events) exposed each event reference without a partition_key.
• The sibling schema used for inlet_events (AssetEventResponse) already carries partition_key under the same unreleased feature, so the consumed-event reference was just an incomplete rollout of it.

What

• Add partition_key to AssetEventDagRunReference (core datamodel) and regenerate the task-sdk _generated.py from it.
• The field is mapper-agnostic: the scheduler reads consumed events from PartitionedAssetKeyLog regardless of mapper kind, so rollup (N→1, multiple events), fan-out (1→N, single event), and identity (1→1, single event) all populate it. Authors read it run-scoped, with no inlet wiring or DB lookup: [e.partition_key for e in triggering_asset_events[asset]].

from __future__ importannotationsfromairflow.sdkimport (
DAG,
Asset,
FixedKeyMapper,
PartitionedAssetTimetable,
PartitionedAtRuntime,
RollupMapper,
SegmentWindow,
task,
)
region_stats=Asset(name="region_stats", uri="memory://region_stats")
withDAG(
dag_id="produce_region_stats",
schedule=PartitionedAtRuntime(),
catchup=False,
tags=["play", "partitions"],
):
@task(outlets=[region_stats])defemit_regions(*, outlet_events):
"""Emit one partition per region from a single run."""outlet_events[region_stats].add_partitions(["us", "eu", "apac"])
emit_regions()
withDAG(
dag_id="rollup_region_stats",
schedule=PartitionedAssetTimetable(
assets=region_stats,
default_partition_mapper=RollupMapper(
upstream_mapper=FixedKeyMapper("all_regions"),
window=SegmentWindow(["us", "eu", "apac"]),
),
),
catchup=False,
tags=["play", "partitions"],
):
@taskdefsummarize(*, dag_run, triggering_asset_events):
"""Print the downstream key and the upstream partitions that produced it."""contributed= [event.partition_keyforeventintriggering_asset_events[region_stats]]
print(f"downstream partition = {dag_run.partition_key}")
print(f"contributing upstream partitions = {sorted(contributed)}")
summarize()

---

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

• Yes (please specify the tool below)

---

• 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.

[goransh-walia]

Nice clean implementation. One minor thought: adding a docstring would help future contributors. But overall this looks solid.

[github-actions]

Backport successfully created: v3-3-test

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
✅	v3-3-test