Reject mismatched rollup mapper and window pairings at Dag parse time

airflow-core/src/airflow/partition_mappers/base.py

@@ -38,10 +38,12 @@ class PartitionMapper(ABC):
 
     is_rollup: ClassVar[bool] = False
 
-    # Class-level default so the attribute resolves even on subclasses whose
-    # __init__ (e.g. attrs-generated in plugins) never calls this base __init__.
     max_downstream_keys: int | None = None
 
+    # Declared result type of decode_downstream; RollupMapper rejects windows
+    # whose expected_decoded_type does not match. Temporal mappers use datetime.
+    expected_decoded_type: ClassVar[type] = str
+
     def __init__(self, *, max_downstream_keys: int | None = None) -> None:
         if max_downstream_keys is not None and (
             not isinstance(max_downstream_keys, int) or max_downstream_keys < 1
@@ -168,16 +170,13 @@ def __init__(
         wait_policy: WaitPolicy | None = None,
         max_downstream_keys: int | None = None,
     ) -> None:
-        decode_overridden = type(upstream_mapper).decode_downstream is not PartitionMapper.decode_downstream
-        if not decode_overridden and window.expected_decoded_type is not str:
+        if upstream_mapper.expected_decoded_type is not window.expected_decoded_type:
             raise TypeError(
                 f"{type(window).__name__} expects decoded values of type "
                 f"{window.expected_decoded_type.__name__!r}, but "
-                f"{type(upstream_mapper).__name__} does not override "
-                f"'decode_downstream' (base default returns str). Pair the window "
-                f"with an upstream mapper that decodes to "
-                f"{window.expected_decoded_type.__name__}, or use a window whose "
-                f"'expected_decoded_type' accepts str."
+                f"{type(upstream_mapper).__name__} decodes to "
+                f"{upstream_mapper.expected_decoded_type.__name__!r}. Pair a window and an "
+                f"upstream mapper whose 'expected_decoded_type' match."
             )
         if wait_policy is None:
             wait_policy = WaitForAll()

airflow-core/src/airflow/partition_mappers/temporal.py

@@ -155,6 +155,8 @@ def _compile_output_format_regex(
 class _BaseTemporalMapper(PartitionMapper):
     """Base class for Temporal Partition Mappers."""
 
+    expected_decoded_type: ClassVar[type] = datetime
+
     default_output_format: str
 
     def __init__(

airflow-core/tests/unit/partition_mappers/test_base.py

@@ -22,9 +22,10 @@
 import pytest
 
 from airflow.partition_mappers.base import PartitionMapper, RollupMapper
+from airflow.partition_mappers.fixed_key import FixedKeyMapper
 from airflow.partition_mappers.identity import IdentityMapper
 from airflow.partition_mappers.temporal import StartOfDayMapper
-from airflow.partition_mappers.window import DayWindow
+from airflow.partition_mappers.window import DayWindow, SegmentWindow
 from airflow.serialization.decoders import decode_partition_mapper
 from airflow.serialization.encoders import encode_partition_mapper
 from airflow.serialization.enums import Encoding
@@ -146,6 +147,69 @@ def to_upstream(self, decoded_downstream):
         # Should not raise.
         RollupMapper(upstream_mapper=_StringOnlyMapper(), window=_AlphaWindow())
 
+    def test_accepts_decode_overriding_str_mapper_with_str_window(self):
+        """Overriding decode/encode while still decoding to str is a legal str-window pairing.
+
+        The guard must key off the declared ``expected_decoded_type`` (str by
+        default), not off whether ``decode_downstream`` is overridden.
+        """
+
+        class _CasefoldMapper(PartitionMapper):
+            def to_downstream(self, key):
+                return key.casefold()
+
+            def decode_downstream(self, downstream_key):
+                return downstream_key.casefold()
+
+            def encode_upstream(self, decoded_upstream):
+                return str(decoded_upstream)
+
+        # Should not raise.
+        RollupMapper(upstream_mapper=_CasefoldMapper(), window=SegmentWindow(["us", "eu"]))
+
+    @pytest.mark.parametrize(
+        ("upstream_mapper_factory", "window_factory", "match"),
+        [
+            pytest.param(
+                lambda: FixedKeyMapper("all"),
+                lambda: SegmentWindow(["us", "eu"]),
+                None,
+                id="str_mapper-str_window-valid",
+            ),
+            pytest.param(
+                StartOfDayMapper,
+                DayWindow,
+                None,
+                id="datetime_mapper-datetime_window-valid",
+            ),
+            pytest.param(
+                lambda: FixedKeyMapper("all"),
+                DayWindow,
+                "DayWindow expects decoded values of type 'datetime'",
+                id="str_mapper-datetime_window-invalid",
+            ),
+            pytest.param(
+                StartOfDayMapper,
+                lambda: SegmentWindow(["us"]),
+                "SegmentWindow expects decoded values of type 'str'",
+                id="datetime_mapper-str_window-invalid",
+            ),
+        ],
+    )
+    def test_upstream_mapper_window_type_pairing(self, upstream_mapper_factory, window_factory, match):
+        """RollupMapper's guard must catch a decoded-type mismatch in either direction.
+
+        Guards the ``datetime_mapper-str_window`` direction in particular: left
+        unchecked, that pairing survives construction and serialization, then
+        raises ``AttributeError`` inside ``to_upstream`` at scheduler tick time.
+        """
+        if match is None:
+            # Should not raise.
+            RollupMapper(upstream_mapper=upstream_mapper_factory(), window=window_factory())
+        else:
+            with pytest.raises(TypeError, match=match):
+                RollupMapper(upstream_mapper=upstream_mapper_factory(), window=window_factory())
+
 
 class TestPartitionMapperMaxDownstreamKeysValidator:
     """Verify the max_downstream_keys validator on the PartitionMapper base class.

task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py

@@ -82,12 +82,11 @@ def __attrs_post_init__(self) -> None:
         # would otherwise be swallowed by the bare ``except`` in
         # ``_create_dagruns_for_partitioned_asset_dags`` and surface only as
         # "Failed to deserialize Dag" spam).
-        if self.upstream_mapper.expected_decoded_type is str and self.window.expected_decoded_type is not str:
+        if self.upstream_mapper.expected_decoded_type is not self.window.expected_decoded_type:
             raise TypeError(
                 f"{type(self.window).__name__} expects decoded values of type "
                 f"{self.window.expected_decoded_type.__name__!r}, but "
-                f"{type(self.upstream_mapper).__name__} decodes to 'str' (SDK PartitionMapper default). "
-                f"Pair the window with an upstream mapper whose 'expected_decoded_type' is "
-                f"{self.window.expected_decoded_type.__name__}, or use a window whose "
-                f"'expected_decoded_type' accepts str."
+                f"{type(self.upstream_mapper).__name__} decodes to "
+                f"{self.upstream_mapper.expected_decoded_type.__name__!r}. Pair a window and an "
+                f"upstream mapper whose 'expected_decoded_type' match."
             )

task-sdk/tests/task_sdk/definitions/test_partition_mappers.py

@@ -192,14 +192,47 @@ def test_rejects_invalid_segments(self, segments, match):
 
 
 class TestSdkCategoricalRollupGuard:
-    """SDK-side RollupMapper guard mirrors core: str mapper + str window passes."""
+    """SDK-side RollupMapper guard mirrors core: the decoded type must match in either direction."""
 
-    def test_fixed_key_with_segment_window_does_not_raise(self):
-        # SDK guard: FixedKeyMapper.expected_decoded_type is str,
-        # SegmentWindow.expected_decoded_type is str -> guard passes.
-        RollupMapper(upstream_mapper=FixedKeyMapper("all"), window=SegmentWindow(["us", "eu"]))
-
-    def test_str_mapper_with_datetime_window_raises(self):
-        # SDK guard: FixedKeyMapper (str) + DayWindow (datetime) -> raise.
-        with pytest.raises(TypeError, match="DayWindow expects decoded values of type 'datetime'"):
-            RollupMapper(upstream_mapper=FixedKeyMapper("all"), window=DayWindow())
+    @pytest.mark.parametrize(
+        ("upstream_mapper_factory", "window_factory", "match"),
+        [
+            pytest.param(
+                lambda: FixedKeyMapper("all"),
+                lambda: SegmentWindow(["us", "eu"]),
+                None,
+                id="str_mapper-str_window-valid",
+            ),
+            pytest.param(
+                StartOfDayMapper,
+                DayWindow,
+                None,
+                id="datetime_mapper-datetime_window-valid",
+            ),
+            pytest.param(
+                lambda: FixedKeyMapper("all"),
+                DayWindow,
+                "DayWindow expects decoded values of type 'datetime'",
+                id="str_mapper-datetime_window-invalid",
+            ),
+            pytest.param(
+                StartOfDayMapper,
+                lambda: SegmentWindow(["us"]),
+                "SegmentWindow expects decoded values of type 'str'",
+                id="datetime_mapper-str_window-invalid",
+            ),
+        ],
+    )
+    def test_upstream_mapper_window_type_pairing(self, upstream_mapper_factory, window_factory, match):
+        """RollupMapper's guard must catch a decoded-type mismatch in either direction.
+
+        Guards the ``datetime_mapper-str_window`` direction in particular: left
+        unchecked, that pairing survives construction and serialization, then
+        raises ``AttributeError`` inside ``to_upstream`` at scheduler tick time.
+        """
+        if match is None:
+            # Should not raise.
+            RollupMapper(upstream_mapper=upstream_mapper_factory(), window=window_factory())
+        else:
+            with pytest.raises(TypeError, match=match):
+                RollupMapper(upstream_mapper=upstream_mapper_factory(), window=window_factory())