This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 79e9628fa1d Add example plugins and expand docs for asset partitions
(#68889)
79e9628fa1d is described below
commit 79e9628fa1d8b35800d1f2fd1b5ed4ddfa0046ab
Author: Wei Lee <[email protected]>
AuthorDate: Fri Jun 26 15:54:03 2026 +0800
Add example plugins and expand docs for asset partitions (#68889)
---
.../docs/authoring-and-scheduling/assets.rst | 267 +++++++++++++++++++++
.../example_dags/example_asset_partition.py | 10 +-
.../plugins/custom_partition_mapper.py | 96 ++++++++
.../plugins/custom_partition_timetable.py | 55 +++++
.../core_api/routes/public/test_plugins.py | 10 +-
.../tests/unit/plugins/test_plugins_manager.py | 4 +-
6 files changed, 434 insertions(+), 8 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst
b/airflow-core/docs/authoring-and-scheduling/assets.rst
index a74664302da..b5561392587 100644
--- a/airflow-core/docs/authoring-and-scheduling/assets.rst
+++ b/airflow-core/docs/authoring-and-scheduling/assets.rst
@@ -523,6 +523,33 @@ creates asset events with a partition key on each run.
Partitioned events are intended for partition-aware downstream scheduling, and
do not trigger non-partition-aware Dags.
+Pre-determined vs runtime partitioning
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Both kinds attach a partition key to a Dag run — the difference is *when* and
+*by whom* the key is decided:
+
+* **Pre-determined partitioning** — the partition key is worked out before the
+ task runs, using the timetable's schedule cadence and partition mappers to
match
+ upstream keys to downstream keys and trigger partition-based Dag runs.
+ :class:`~airflow.sdk.CronPartitionTimetable` uses this kind as a producer;
+ :class:`~airflow.sdk.PartitionedAssetTimetable` uses it as a consumer.
+
+* **Runtime partitioning** — the partition key is deferred to task runtime:
+ the producing task records key(s) via
``outlet_events[self].add_partitions(...)``.
+ :class:`~airflow.sdk.PartitionedAtRuntime` uses this kind and never
schedules on its
+ own (``can_be_scheduled=False``); a schedulable timetable can also defer to
runtime
+ by subclassing :class:`~airflow.timetables.trigger.CronTriggerTimetable` and
setting
+ ``partitioned_at_runtime = True`` (see the custom plugin example below).
+
+A timetable uses one kind or the other, not both: it either resolves partitions
+ahead of the run or defers them to task runtime.
+
+**Practical rule:** use :class:`~airflow.sdk.CronPartitionTimetable` when the
+partition key follows from the schedule cadence; use
:class:`~airflow.sdk.PartitionedAtRuntime`
+when the key is only known once the task runs (e.g. a watermark from source
data);
+use :class:`~airflow.sdk.PartitionedAssetTimetable` downstream to consume
either kind.
+
For downstream partition-aware scheduling, use ``PartitionedAssetTimetable``:
.. code-block:: python
@@ -753,6 +780,62 @@ so the run is held indefinitely) and the fall-back day has
twenty-five (the repe
hour is dropped). Use a UTC-based upstream mapper for any rollup that crosses
a DST
boundary; see the ``DayWindow`` class docstring for the full discussion.
+Wait policies
+~~~~~~~~~~~~~
+
+:class:`~airflow.sdk.RollupMapper` accepts an optional ``wait_policy`` argument
+that decides when the downstream Dag run fires given the expected window vs the
+upstream keys that have actually arrived.
+
+* :class:`~airflow.sdk.WaitForAll` (the default) holds the run until every
expected
+ upstream key in the window has arrived.
+* :class:`~airflow.sdk.MinimumCount` ``(n)`` fires early once at least ``n``
of the
+ expected keys have arrived — useful to tolerate a slow or missing upstream
partition
+ rather than holding the run indefinitely.
+
+.. code-block:: python
+
+ from airflow.sdk import (
+ DAG,
+ Asset,
+ FixedKeyMapper,
+ MinimumCount,
+ PartitionedAtRuntime,
+ PartitionedAssetTimetable,
+ RollupMapper,
+ SegmentWindow,
+ asset,
+ )
+
+
+ @asset(
+ uri="file://incoming/player-stats/multi-region.csv",
+ schedule=PartitionedAtRuntime(),
+ )
+ def multi_region_player_stats(self, outlet_events):
+ outlet_events[self].add_partitions(["us", "eu", "apac"])
+
+
+ # Consumer: fires once at least two of the three declared region
partitions arrive.
+ with DAG(
+ dag_id="segment_region_stats_early_rollup",
+ schedule=PartitionedAssetTimetable(
+ assets=Asset.ref(name="multi_region_player_stats"),
+ default_partition_mapper=RollupMapper(
+ upstream_mapper=FixedKeyMapper("all_regions"),
+ window=SegmentWindow(["us", "eu", "apac"]),
+ wait_policy=MinimumCount(2),
+ ),
+ ),
+ catchup=False,
+ ):
+ ...
+
+``MinimumCount(-1)`` is the relative spelling of the same threshold — "at most
one
+missing" — and is equivalent to ``MinimumCount(2)`` for a three-member window.
+Pass :class:`~airflow.sdk.WaitForAll` explicitly when you want to document
intent
+rather than relying on the default.
+
.. _segment-categorical-rollup:
Segment (categorical) rollup
@@ -872,5 +955,189 @@ When a runtime run emits exactly one partition key, the
producing
these events the same way as timetable-produced partitions, through
``PartitionedAssetTimetable``.
+Fan-out mappers
+~~~~~~~~~~~~~~~
+
+.. versionadded:: 3.3.0
+
+:class:`~airflow.sdk.FanOutMapper` is the mirror of
:class:`~airflow.sdk.RollupMapper`:
+instead of holding many upstream events until one downstream run can fire, a
single
+upstream event fans *out* to one downstream Dag run per window member. It
composes an
+``upstream_mapper`` (which normalizes the upstream key to the window anchor)
with a
+:class:`~airflow.sdk.Window` that enumerates the downstream period, and an
optional
+``downstream_mapper`` that converts each window member into a downstream
partition key
+string.
+
+For temporal windows (:class:`~airflow.sdk.WeekWindow`,
:class:`~airflow.sdk.MonthWindow`,
+etc.) a default ``downstream_mapper`` is applied automatically — for example
+:class:`~airflow.sdk.WeekWindow` defaults to
:class:`~airflow.sdk.StartOfDayMapper`,
+so each of the seven daily members is encoded as a ``YYYY-MM-DD`` string. For
+:class:`~airflow.sdk.SegmentWindow` there is no default-table entry, so an
explicit
+``downstream_mapper`` is required.
+
+The following example fans a weekly model artifact out to seven daily
inference runs —
+one Dag run per day in the week:
+
+.. code-block:: python
+
+ from airflow.sdk import (
+ DAG,
+ Asset,
+ CronPartitionTimetable,
+ FanOutMapper,
+ PartitionedAssetTimetable,
+ StartOfWeekMapper,
+ WeekWindow,
+ task,
+ )
+
+ weekly_model_artifact = Asset(uri="file://artifacts/models/weekly.bin",
name="weekly_model_artifact")
+
+ # Producer: emits one partitioned event per week (key is the Monday date).
+ with DAG(
+ dag_id="train_weekly_model",
+ schedule=CronPartitionTimetable("0 0 * * 1", timezone="UTC"),
+ catchup=False,
+ ):
+
+ @task(outlets=[weekly_model_artifact])
+ def train_model():
+ pass
+
+ train_model()
+
+
+ # Consumer: one Dag run per day derived from the weekly upstream event.
+ with DAG(
+ dag_id="daily_inference",
+ schedule=PartitionedAssetTimetable(
+ assets=weekly_model_artifact,
+ default_partition_mapper=FanOutMapper(
+ upstream_mapper=StartOfWeekMapper(),
+ window=WeekWindow(),
+ max_downstream_keys=7,
+ ),
+ ),
+ catchup=False,
+ ):
+
+ @task
+ def run_inference(dag_run=None):
+ # dag_run.partition_key is one daily key, e.g. "2026-03-10".
+ print(dag_run.partition_key)
+
+ run_inference()
+
+``max_downstream_keys`` caps how many downstream Dag runs one upstream event
may
+create. When exceeded, the runs for that event are **not** queued and a
"partition
+fan-out exceeded" audit-log entry is recorded instead. Omitting it falls back
to the
+global ``[scheduler] partition_mapper_max_downstream_keys`` config (default
1000).
+Set it explicitly to document intent and guard against accidental fan-out
explosions.
+
+Window direction: FORWARD and BACKWARD
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Every :class:`~airflow.sdk.Window` supports a ``direction`` parameter that
controls
+which period the window enumerates relative to its anchor.
+
+* ``Window.Direction.FORWARD`` (the default) — yields the period *starting* at
the
+ upstream key. For a weekly upstream key ``"2026-03-09"`` (Monday),
+ ``WeekWindow()`` yields the seven days ``2026-03-09`` through ``2026-03-15``.
+* ``Window.Direction.BACKWARD`` — yields the trailing period *ending* at the
key.
+ The same ``"2026-03-09"`` key with
``WeekWindow(direction=Window.Direction.BACKWARD)``
+ yields the seven days ending on that Monday (``2026-03-03`` through
``2026-03-09``).
+
+.. code-block:: python
+
+ from airflow.sdk import FanOutMapper, PartitionedAssetTimetable,
StartOfWeekMapper, WeekWindow, Window
+
+ PartitionedAssetTimetable(
+ assets=weekly_model_artifact,
+ default_partition_mapper=FanOutMapper(
+ upstream_mapper=StartOfWeekMapper(),
+ window=WeekWindow(direction=Window.Direction.BACKWARD),
+ ),
+ )
+
+Direction applies to rollup windows too — ``RollupMapper`` uses the same window
+classes, so ``DayWindow(direction=Window.Direction.BACKWARD)`` holds a
downstream
+run until the twenty-four hours *preceding* midnight of the downstream key have
+all arrived.
+
+Custom partition mappers, windows, and timetables in plugins
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Custom :class:`~airflow.partition_mappers.base.PartitionMapper`,
+:class:`~airflow.partition_mappers.window.Window`, and partition-aware
+:class:`Timetable <airflow.timetables.base.Timetable>` classes can be shipped
as
+Airflow plugins by listing them in ``AirflowPlugin.partition_mappers``,
+``.windows``, and ``.timetables`` respectively. Once a plugin is installed,
+these classes become usable in :class:`~airflow.sdk.PartitionedAssetTimetable`
+and :class:`~airflow.partition_mappers.base.RollupMapper` without modifying
core
+Airflow.
+
+**Custom partition mapper** — strip a namespace prefix so that upstream keys
+like ``"eu::daily-sales"`` and ``"us::daily-sales"`` both collapse to the
+downstream key ``"daily-sales"``:
+
+.. exampleinclude::
/../src/airflow/example_dags/plugins/custom_partition_mapper.py
+ :language: python
+ :start-after: [START custom_partition_mapper]
+ :end-before: [END custom_partition_mapper]
+
+**Custom rollup window** — yield only weekday period-starts from a calendar
+month so a downstream asset waits only for business-day upstream partitions:
+
+.. exampleinclude:: /../src/airflow/example_dags/plugins/business_day_window.py
+ :language: python
+ :start-after: [START custom_window]
+ :end-before: [END custom_window]
+
+**Custom runtime-partitioned timetable** — a schedulable cron timetable that
+defers the partition key to task runtime, so the producing task can check
whether
+the period's data exists before emitting a partition:
+
+.. exampleinclude::
/../src/airflow/example_dags/plugins/custom_partition_timetable.py
+ :language: python
+ :start-after: [START custom_partition_timetable]
+ :end-before: [END custom_partition_timetable]
+
+A producer Dag schedules on the timetable's cron cadence and decides the
+partition key at runtime — emitting it only when the period's upstream data is
+present. If the data has not arrived, the task simply does not call
+``add_partitions`` (an empty or ``None`` key is rejected anyway), so no
+partitioned event — and therefore no downstream ``PartitionedAssetTimetable``
+run — is produced for that period:
+
+.. code-block:: python
+
+ from airflow.sdk import DAG, Asset, task
+
+ # ScheduledRuntimePartitionTimetable is provided by the plugin registered
above.
+ from my_plugin.plugins import ScheduledRuntimePartitionTimetable
+
+ daily_export = Asset(uri="file://exports/daily.csv", name="daily_export")
+
+ with DAG(
+ dag_id="export_when_ready",
+ schedule=ScheduledRuntimePartitionTimetable("0 6 * * *",
timezone="UTC"),
+ catchup=False,
+ ):
+
+ @task(outlets=[daily_export])
+ def export(*, outlet_events):
+ # Fires every day at 06:00 UTC. The partition key is not fixed by
the
+ # schedule — decide it at runtime from the data itself: find which
+ # day's source file has actually landed.
+ partition_key = latest_ready_day("s3://raw") # your own check;
e.g. "2026-06-23" or None
+ if partition_key:
+ build_export(partition_key)
+ outlet_events[daily_export].add_partitions(partition_key)
+ # If nothing is ready, emit nothing: with no partition key
recorded,
+ # no partitioned event is produced and no downstream
+ # PartitionedAssetTimetable run is triggered for this period.
+
+ export()
+
For complete runnable examples, see
``airflow-core/src/airflow/example_dags/example_asset_partition.py``.
diff --git a/airflow-core/src/airflow/example_dags/example_asset_partition.py
b/airflow-core/src/airflow/example_dags/example_asset_partition.py
index f4171bb46f1..e7151680404 100644
--- a/airflow-core/src/airflow/example_dags/example_asset_partition.py
+++ b/airflow-core/src/airflow/example_dags/example_asset_partition.py
@@ -44,6 +44,7 @@ from airflow.sdk import (
WeekWindow,
Window,
asset,
+ get_current_context,
task,
)
@@ -217,9 +218,14 @@ with DAG(
"""
@task(outlets=[region_raw_stats])
- def ingest_region():
+ def ingest_region(dag_run=None):
"""Materialize player statistics for a single region partition."""
- pass
+ context = get_current_context()
+ if TYPE_CHECKING:
+ assert dag_run
+ print(
+ f"dag_run partition key {dag_run.partition_key} context partition
key {context['partition_key']}"
+ )
ingest_region()
diff --git
a/airflow-core/src/airflow/example_dags/plugins/custom_partition_mapper.py
b/airflow-core/src/airflow/example_dags/plugins/custom_partition_mapper.py
new file mode 100644
index 00000000000..ab6ff673044
--- /dev/null
+++ b/airflow-core/src/airflow/example_dags/plugins/custom_partition_mapper.py
@@ -0,0 +1,96 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.partition_mappers.base import PartitionMapper
+from airflow.plugins_manager import AirflowPlugin
+
+if TYPE_CHECKING:
+ from typing import Any
+
+
+# [START custom_partition_mapper]
+class PrefixStripMapper(PartitionMapper):
+ """
+ A partition mapper that strips a fixed namespace prefix from upstream keys.
+
+ Upstream systems often qualify partition keys with a region or environment
+ prefix — for example ``"eu::daily-sales"`` or ``"us::daily-sales"``. A
+ downstream asset that aggregates across regions only cares about the base
key
+ (``"daily-sales"``). ``PrefixStripMapper`` strips the given prefix
(including
+ a configurable separator) so that all upstream namespaces collapse to the
+ same downstream partition key.
+
+ If the upstream key does not start with the configured prefix the key is
+ returned unchanged, which is deliberate: keys that already live in the
target
+ namespace pass through without modification.
+
+ This class demonstrates registering a custom :class:`PartitionMapper
+ <airflow.partition_mappers.base.PartitionMapper>` subclass via the
+ ``AirflowPlugin.partition_mappers`` registry. Any plugin that lists it in
+ ``partition_mappers = [...]`` makes it available to
+ :class:`~airflow.sdk.PartitionedAssetTimetable` and
+ :class:`~airflow.partition_mappers.base.RollupMapper` without modifying
core
+ Airflow.
+
+ :param prefix: The namespace prefix to strip, e.g. ``"eu"``.
+ :param separator: The string that separates the prefix from the base key.
+ Defaults to ``"::"`` to match a common ``"region::key"`` convention.
+ """
+
+ def __init__(
+ self,
+ prefix: str,
+ *,
+ separator: str = "::",
+ max_downstream_keys: int | None = None,
+ ) -> None:
+ super().__init__(max_downstream_keys=max_downstream_keys)
+ if not prefix:
+ raise ValueError("prefix must be a non-empty string.")
+ self.prefix = prefix
+ self.separator = separator
+
+ def to_downstream(self, key: str) -> str:
+ full_prefix = self.prefix + self.separator
+ if key.startswith(full_prefix):
+ return key[len(full_prefix) :]
+ return key
+
+ def serialize(self) -> dict[str, Any]:
+ data: dict[str, Any] = {"prefix": self.prefix, "separator":
self.separator}
+ if self.max_downstream_keys is not None:
+ data["max_downstream_keys"] = self.max_downstream_keys
+ return data
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> PartitionMapper:
+ return cls(
+ prefix=data["prefix"],
+ separator=data.get("separator", "::"),
+ max_downstream_keys=data.get("max_downstream_keys"),
+ )
+
+
+class PrefixStripMapperPlugin(AirflowPlugin):
+ name = "prefix_strip_mapper_plugin"
+ partition_mappers = [PrefixStripMapper]
+
+
+# [END custom_partition_mapper]
diff --git
a/airflow-core/src/airflow/example_dags/plugins/custom_partition_timetable.py
b/airflow-core/src/airflow/example_dags/plugins/custom_partition_timetable.py
new file mode 100644
index 00000000000..f74bc99b62e
--- /dev/null
+++
b/airflow-core/src/airflow/example_dags/plugins/custom_partition_timetable.py
@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow.plugins_manager import AirflowPlugin
+from airflow.timetables.trigger import CronTriggerTimetable
+
+
+# [START custom_partition_timetable]
+class ScheduledRuntimePartitionTimetable(CronTriggerTimetable):
+ """
+ A schedulable timetable whose partition key is decided at task runtime.
+
+ Runs fire on the given cron cadence, exactly like an ordinary
+ :class:`~airflow.timetables.trigger.CronTriggerTimetable`. The partition
key,
+ however, is not derived from the schedule: it is set while the producing
task
+ runs — typically after the task checks whether the period's source data has
+ arrived — by calling ``outlet_events[self].add_partitions(...)``.
+
+ This uses runtime partitioning on a regular cron schedule: the timetable
stays
+ schedulable (``can_be_scheduled`` is ``True``) yet sets
+ ``partitioned_at_runtime = True`` so the partition key is deferred to task
+ runtime. It differs from :class:`~airflow.sdk.PartitionedAtRuntime` (which
also
+ defers the key to runtime but never schedules a run on its own) and from
+ :class:`~airflow.timetables.trigger.CronPartitionTimetable` (which works
out the
+ partition key from the cadence ahead of the run). ``partitioned`` stays
+ ``False``: no partition key is worked out ahead of the run.
+
+ Registering it via the ``AirflowPlugin.timetables`` registry makes it
usable
+ by Dag authors without modifying core Airflow.
+ """
+
+ partitioned_at_runtime = True
+
+
+class ScheduledRuntimePartitionTimetablePlugin(AirflowPlugin):
+ name = "scheduled_runtime_partition_timetable_plugin"
+ timetables = [ScheduledRuntimePartitionTimetable]
+
+
+# [END custom_partition_timetable]
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
index 42df1955ba5..6bca49d8052 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
@@ -34,7 +34,7 @@ class TestGetPlugins:
# Filters
(
{},
- 16,
+ 18,
[
"InformaticaProviderPlugin",
"MetadataCollectionPlugin",
@@ -49,21 +49,23 @@ class TestGetPlugins:
"plugin-b",
"plugin-c",
"postload",
+ "prefix_strip_mapper_plugin",
"priority_weight_strategy_plugin",
+ "scheduled_runtime_partition_timetable_plugin",
"test_plugin",
"workday_timetable_plugin",
],
),
(
{"limit": 3, "offset": 3},
- 16,
+ 18,
[
"business_day_window_plugin",
"databricks_workflow",
"decreasing_priority_weight_strategy_plugin",
],
),
- ({"limit": 1}, 16, ["InformaticaProviderPlugin"]),
+ ({"limit": 1}, 18, ["InformaticaProviderPlugin"]),
],
)
def test_should_respond_200(
@@ -166,7 +168,7 @@ class TestGetPlugins:
assert len(plugins_page) == 7
assert "test_plugin_invalid" not in [p["name"] for p in plugins_page]
- assert body["total_entries"] == 16
+ assert body["total_entries"] == 18
@skip_if_force_lowest_dependencies_marker
diff --git a/airflow-core/tests/unit/plugins/test_plugins_manager.py
b/airflow-core/tests/unit/plugins/test_plugins_manager.py
index 4970b180ad2..ce55db1e269 100644
--- a/airflow-core/tests/unit/plugins/test_plugins_manager.py
+++ b/airflow-core/tests/unit/plugins/test_plugins_manager.py
@@ -80,7 +80,7 @@ class TestPluginsManager:
example_plugins_module="airflow.example_dags.plugins",
)
- assert len(plugins) == 11
+ assert len(plugins) == 13
assert not import_errors
for plugin in plugins:
if "AirflowTestOnLoadPlugin" in str(plugin):
@@ -103,7 +103,7 @@ class TestPluginsManager:
):
plugins, import_errors = plugins_manager._get_plugins()
- assert len(plugins) == 4 # four are loaded from examples
+ assert len(plugins) == 6 # four are loaded from examples
assert len(import_errors) == 1
received_logs = caplog.text
