feat: make temporal partition_mapper timezone aware

Lee-W · Lee-W · commit 3f8d15b2ad42 · 2026-03-09T16:58:38.000+08:00
diff --git a/airflow-core/src/airflow/example_dags/example_asset_partition.py b/airflow-core/src/airflow/example_dags/example_asset_partition.py
@@ -77,7 +77,7 @@ def team_c_player_stats():
     dag_id="clean_and_combine_player_stats",
     schedule=PartitionedAssetTimetable(
         assets=team_a_player_stats & team_b_player_stats & team_c_player_stats,
-        default_partition_mapper=HourlyMapper(),
+        default_partition_mapper=HourlyMapper(timezone="UTC"),
     ),
     catchup=False,
     tags=["player-stats", "cleanup"],
@@ -121,9 +121,9 @@ def compute_player_odds():
     schedule=PartitionedAssetTimetable(
         assets=(combined_player_stats & team_a_player_stats & Asset.ref(name="team_b_player_stats")),
         partition_mapper_config={
-            combined_player_stats: YearlyMapper(),  # incompatible on purpose
-            team_a_player_stats: HourlyMapper(),
-            Asset.ref(name="team_b_player_stats"): HourlyMapper(),
+            combined_player_stats: YearlyMapper(timezone="UTC"),  # incompatible on purpose
+            team_a_player_stats: HourlyMapper(timezone="UTC"),
+            Asset.ref(name="team_b_player_stats"): HourlyMapper(timezone="UTC"),
         },
     ),
     catchup=False,
diff --git a/airflow-core/src/airflow/partition_mappers/temporal.py b/airflow-core/src/airflow/partition_mappers/temporal.py
@@ -18,10 +18,14 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
+from airflow._shared.timezones.timezone import parse_timezone
 from airflow.partition_mappers.base import PartitionMapper
 
+if TYPE_CHECKING:
+    from pendulum import FixedTimezone, Timezone
+
 
 class _BaseTemporalMapper(PartitionMapper, ABC):
     """Base class for Temporal Partition Mappers."""
@@ -30,14 +34,24 @@ class _BaseTemporalMapper(PartitionMapper, ABC):
 
     def __init__(
         self,
-        input_format: str = "%Y-%m-%dT%H:%M:%S",
+        *,
+        timezone: str | Timezone | FixedTimezone,
+        input_format: str = "%Y-%m-%dT%H:%M:%S%z",
         output_format: str | None = None,
     ):
         self.input_format = input_format
         self.output_format = output_format or self.default_output_format
+        if isinstance(timezone, str):
+            timezone = parse_timezone(timezone)
+        self._timezone: Timezone | FixedTimezone = timezone
 
     def to_downstream(self, key: str) -> str:
         dt = datetime.strptime(key, self.input_format)
+        if dt.tzinfo is None:
+            dt = dt.replace(tzinfo=self._timezone)
+        else:
+            dt = dt.astimezone(self._timezone)
+
         normalized = self.normalize(dt)
         return self.format(normalized)
 
@@ -50,14 +64,18 @@ def format(self, dt: datetime) -> str:
         return dt.strftime(self.output_format)
 
     def serialize(self) -> dict[str, Any]:
+        from airflow.serialization.encoders import encode_timezone
+
         return {
+            "timezone": encode_timezone(self._timezone),
             "input_format": self.input_format,
             "output_format": self.output_format,
         }
 
     @classmethod
     def deserialize(cls, data: dict[str, Any]) -> PartitionMapper:
         return cls(
+            timezone=parse_timezone(data["timezone"]),
             input_format=data["input_format"],
             output_format=data["output_format"],
         )
@@ -66,7 +84,7 @@ def deserialize(cls, data: dict[str, Any]) -> PartitionMapper:
 class HourlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to hour."""
 
-    default_output_format = "%Y-%m-%dT%H"
+    default_output_format = "%Y-%m-%dT%H%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(minute=0, second=0, microsecond=0)
@@ -75,7 +93,7 @@ def normalize(self, dt: datetime) -> datetime:
 class DailyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to day."""
 
-    default_output_format = "%Y-%m-%d"
+    default_output_format = "%Y-%m-%d%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(hour=0, minute=0, second=0, microsecond=0)
@@ -84,7 +102,7 @@ def normalize(self, dt: datetime) -> datetime:
 class WeeklyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to week."""
 
-    default_output_format = "%Y-%m-%d (W%V)"
+    default_output_format = "%Y-%m-%d (W%V)%z"
 
     def normalize(self, dt: datetime) -> datetime:
         start = dt - timedelta(days=dt.weekday())
@@ -94,7 +112,7 @@ def normalize(self, dt: datetime) -> datetime:
 class MonthlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to month."""
 
-    default_output_format = "%Y-%m"
+    default_output_format = "%Y-%m%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(
@@ -109,7 +127,7 @@ def normalize(self, dt: datetime) -> datetime:
 class QuarterlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to quarter."""
 
-    default_output_format = "%Y-Q{quarter}"
+    default_output_format = "%Y-Q{quarter}%z"
 
     def normalize(self, dt: datetime) -> datetime:
         quarter = (dt.month - 1) // 3
@@ -131,7 +149,7 @@ def format(self, dt: datetime) -> str:
 class YearlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to year."""
 
-    default_output_format = "%Y"
+    default_output_format = "%Y%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(
diff --git a/airflow-core/src/airflow/serialization/encoders.py b/airflow-core/src/airflow/serialization/encoders.py
@@ -407,6 +407,7 @@ def _(
         | YearlyMapper,
     ) -> dict[str, Any]:
         return {
+            "timezone": partition_mapper.timezone,
             "input_format": partition_mapper.input_format,
             "output_format": partition_mapper.output_format,
         }
diff --git a/airflow-core/src/airflow/timetables/trigger.py b/airflow-core/src/airflow/timetables/trigger.py
@@ -401,7 +401,7 @@ def __init__(
         run_offset: int | datetime.timedelta | relativedelta | None = None,
         run_immediately: bool | datetime.timedelta = False,
         # todo: AIP-76 we can't infer partition date from this, so we need to store it separately.
-        key_format: str = r"%Y-%m-%dT%H:%M:%S",
+        key_format: str = r"%Y-%m-%dT%H:%M:%S%z",
     ) -> None:
         super().__init__(cron, timezone=timezone, run_immediately=run_immediately)
         if not isinstance(run_offset, (int, NoneType)):
diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py b/airflow-core/tests/unit/cli/commands/test_dag_command.py
@@ -220,11 +220,11 @@ def test_show_dag_imgcat(self, mock_render_dag, mock_popen, stdout_capture):
                 "timedelta(days=-5)",
                 "CronPartitionTimetable('0 0 * * *', timezone='UTC')",
                 "False",
-                jan_1.strftime(r"%Y-%m-%dT%H:%M:%S"),
+                jan_1.strftime(r"%Y-%m-%dT%H:%M:%S%z"),
                 os.linesep.join(
                     [
-                        jan_1.strftime(r"%Y-%m-%dT%H:%M:%S"),
-                        (jan_1 + timedelta(days=1)).strftime(r"%Y-%m-%dT%H:%M:%S"),
+                        jan_1.strftime(r"%Y-%m-%dT%H:%M:%S%z"),
+                        (jan_1 + timedelta(days=1)).strftime(r"%Y-%m-%dT%H:%M:%S%z"),
                     ],
                 ),
                 id="partitioned",
diff --git a/airflow-core/tests/unit/partition_mappers/test_temporal.py b/airflow-core/tests/unit/partition_mappers/test_temporal.py
@@ -33,38 +33,39 @@ class TestTemporalMappers:
     @pytest.mark.parametrize(
         ("mapper_cls", "expected_downstream_key"),
         [
-            (HourlyMapper, "2026-02-10T14"),
-            (DailyMapper, "2026-02-10"),
-            (WeeklyMapper, "2026-02-09 (W07)"),
-            (MonthlyMapper, "2026-02"),
-            (QuarterlyMapper, "2026-Q1"),
-            (YearlyMapper, "2026"),
+            (HourlyMapper, "2026-02-10T14+0000"),
+            (DailyMapper, "2026-02-10+0000"),
+            (WeeklyMapper, "2026-02-09 (W07)+0000"),
+            (MonthlyMapper, "2026-02+0000"),
+            (QuarterlyMapper, "2026-Q1+0000"),
+            (YearlyMapper, "2026+0000"),
         ],
     )
     def test_to_downstream(
         self,
         mapper_cls: type[_BaseTemporalMapper],
         expected_downstream_key: str,
     ):
-        pm = mapper_cls()
-        assert pm.to_downstream("2026-02-10T14:30:45") == expected_downstream_key
+        pm = mapper_cls(timezone="UTC")
+        assert pm.to_downstream("2026-02-10T14:30:45+0000") == expected_downstream_key
 
     @pytest.mark.parametrize(
         ("mapper_cls", "expected_outut_format"),
         [
-            (HourlyMapper, "%Y-%m-%dT%H"),
-            (DailyMapper, "%Y-%m-%d"),
-            (WeeklyMapper, "%Y-%m-%d (W%V)"),
-            (MonthlyMapper, "%Y-%m"),
-            (QuarterlyMapper, "%Y-Q{quarter}"),
-            (YearlyMapper, "%Y"),
+            (HourlyMapper, "%Y-%m-%dT%H%z"),
+            (DailyMapper, "%Y-%m-%d%z"),
+            (WeeklyMapper, "%Y-%m-%d (W%V)%z"),
+            (MonthlyMapper, "%Y-%m%z"),
+            (QuarterlyMapper, "%Y-Q{quarter}%z"),
+            (YearlyMapper, "%Y%z"),
         ],
     )
     def test_serialize(self, mapper_cls: type[_BaseTemporalMapper], expected_outut_format: str):
-        pm = mapper_cls()
+        pm = mapper_cls(timezone="UTC")
         assert pm.serialize() == {
-            "input_format": "%Y-%m-%dT%H:%M:%S",
+            "input_format": "%Y-%m-%dT%H:%M:%S%z",
             "output_format": expected_outut_format,
+            "timezone": "UTC",
         }
 
     @pytest.mark.parametrize(
@@ -73,11 +74,8 @@ def test_serialize(self, mapper_cls: type[_BaseTemporalMapper], expected_outut_f
     )
     def test_deserialize(self, mapper_cls):
         pm = mapper_cls.deserialize(
-            {
-                "input_format": "%Y-%m-%dT%H:%M:%S",
-                "output_format": "customized-format",
-            }
+            {"input_format": "%Y-%m-%dT%H:%M:%S%z", "output_format": "customized-format", "timezone": "UTC"}
         )
         assert isinstance(pm, mapper_cls)
-        assert pm.input_format == "%Y-%m-%dT%H:%M:%S"
+        assert pm.input_format == "%Y-%m-%dT%H:%M:%S%z"
         assert pm.output_format == "customized-format"
diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py b/airflow-core/tests/unit/serialization/test_serialized_objects.py
@@ -751,39 +751,39 @@ def test_encode_timezone():
         (IdentityMapper, [], "airflow.partition_mappers.identity.IdentityMapper", {}),
         (
             HourlyMapper,
-            [],
+            ["UTC"],
             "airflow.partition_mappers.temporal.HourlyMapper",
-            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m-%dT%H"},
+            {"input_format": "%Y-%m-%dT%H:%M:%S%z", "output_format": "%Y-%m-%dT%H%z", "timezone": "UTC"},
         ),
         (
             DailyMapper,
-            [],
+            ["UTC"],
             "airflow.partition_mappers.temporal.DailyMapper",
-            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m-%d"},
+            {"input_format": "%Y-%m-%dT%H:%M:%S%z", "output_format": "%Y-%m-%d%z", "timezone": "UTC"},
         ),
         (
             WeeklyMapper,
-            [],
+            ["UTC"],
             "airflow.partition_mappers.temporal.WeeklyMapper",
-            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m-%d (W%V)"},
+            {"input_format": "%Y-%m-%dT%H:%M:%S%z", "output_format": "%Y-%m-%d (W%V)%z", "timezone": "UTC"},
         ),
         (
             MonthlyMapper,
-            [],
+            ["UTC"],
             "airflow.partition_mappers.temporal.MonthlyMapper",
-            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m"},
+            {"input_format": "%Y-%m-%dT%H:%M:%S%z", "output_format": "%Y-%m%z", "timezone": "UTC"},
         ),
         (
             QuarterlyMapper,
-            [],
+            ["UTC"],
             "airflow.partition_mappers.temporal.QuarterlyMapper",
-            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-Q{quarter}"},
+            {"input_format": "%Y-%m-%dT%H:%M:%S%z", "output_format": "%Y-Q{quarter}%z", "timezone": "UTC"},
         ),
         (
             YearlyMapper,
-            [],
+            ["UTC"],
             "airflow.partition_mappers.temporal.YearlyMapper",
-            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y"},
+            {"input_format": "%Y-%m-%dT%H:%M:%S%z", "output_format": "%Y%z", "timezone": "UTC"},
         ),
     ],
 )
@@ -798,22 +798,22 @@ def test_encode_partition_mapper(cls, args, encode_type, encode_var):
 
 
 @pytest.mark.parametrize(
-    ("sdk_cls", "core_cls"),
+    ("sdk_cls", "args", "core_cls"),
     [
-        (IdentityMapper, CoreIdentityMapper),
-        (HourlyMapper, CoureHourlyMapper),
-        (DailyMapper, CoreDailyMapper),
-        (WeeklyMapper, CoreWeeklyMapper),
-        (MonthlyMapper, CoreMonthlyMapper),
-        (QuarterlyMapper, CoreQuarterlyMapper),
-        (YearlyMapper, CoreYearlyMapper),
+        (IdentityMapper, [], CoreIdentityMapper),
+        (HourlyMapper, ["UTC"], CoureHourlyMapper),
+        (DailyMapper, ["UTC"], CoreDailyMapper),
+        (WeeklyMapper, ["UTC"], CoreWeeklyMapper),
+        (MonthlyMapper, ["UTC"], CoreMonthlyMapper),
+        (QuarterlyMapper, ["UTC"], CoreQuarterlyMapper),
+        (YearlyMapper, ["UTC"], CoreYearlyMapper),
     ],
 )
-def test_decode_partition_mapper(sdk_cls, core_cls):
+def test_decode_partition_mapper(sdk_cls, args, core_cls):
     from airflow.serialization.decoders import decode_partition_mapper
     from airflow.serialization.encoders import encode_partition_mapper
 
-    partition_mapper = sdk_cls()
+    partition_mapper = sdk_cls(*args)
     encoded_pm = encode_partition_mapper(partition_mapper)
 
     core_pm = decode_partition_mapper(encoded_pm)
diff --git a/task-sdk/src/airflow/sdk/definitions/partition_mappers/temporal.py b/task-sdk/src/airflow/sdk/definitions/partition_mappers/temporal.py
@@ -16,52 +16,59 @@
 # under the License.
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from airflow.sdk.definitions.partition_mappers.base import PartitionMapper
 
+if TYPE_CHECKING:
+    from pendulum.tz.timezone import FixedTimezone, Timezone
+
 
 class _BaseTemporalMapper(PartitionMapper):
     default_output_format: str
 
     def __init__(
         self,
-        input_format: str = "%Y-%m-%dT%H:%M:%S",
+        timezone: str | Timezone | FixedTimezone,
+        input_format: str = "%Y-%m-%dT%H:%M:%S%z",
         output_format: str | None = None,
     ) -> None:
+        self.timezone = timezone
         self.input_format = input_format
         self.output_format = output_format or self.default_output_format
 
 
 class HourlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to hour."""
 
-    default_output_format = "%Y-%m-%dT%H"
+    default_output_format = "%Y-%m-%dT%H%z"
 
 
 class DailyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to day."""
 
-    default_output_format = "%Y-%m-%d"
+    default_output_format = "%Y-%m-%d%z"
 
 
 class WeeklyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to week."""
 
-    default_output_format = "%Y-%m-%d (W%V)"
+    default_output_format = "%Y-%m-%d (W%V)%z"
 
 
 class MonthlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to month."""
 
-    default_output_format = "%Y-%m"
+    default_output_format = "%Y-%m%z"
 
 
 class QuarterlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to quarter."""
 
-    default_output_format = "%Y-Q{quarter}"
+    default_output_format = "%Y-Q{quarter}%z"
 
 
 class YearlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to year."""
 
-    default_output_format = "%Y"
+    default_output_format = "%Y%z"
diff --git a/task-sdk/src/airflow/sdk/definitions/timetables/trigger.py b/task-sdk/src/airflow/sdk/definitions/timetables/trigger.py
@@ -168,7 +168,7 @@ def __init__(
         run_offset: int | datetime.timedelta | relativedelta | None = None,
         run_immediately: bool | datetime.timedelta = False,
         # todo: AIP-76 we can't infer partition date from this, so we need to store it separately
-        key_format: str = r"%Y-%m-%dT%H:%M:%S",
+        key_format: str = r"%Y-%m-%dT%H:%M:%S%z",
     ) -> None:
         if not isinstance(run_offset, (int, NoneType)):
             # todo: AIP-76 implement timedelta / relative delta?

Original file line number	Diff line number	Diff line change
`@@ -407,6 +407,7 @@ def _(`
`407`	`407`	`\| YearlyMapper,`
`408`	`408`	`) -> dict[str, Any]:`
`409`	`409`	`return {`
	`410`	`+ "timezone": partition_mapper.timezone,`
`410`	`411`	`"input_format": partition_mapper.input_format,`
`411`	`412`	`"output_format": partition_mapper.output_format,`
`412`	`413`	`}`