|
18 | 18 |
|
19 | 19 | import re |
20 | 20 | from abc import abstractmethod |
| 21 | +from collections.abc import Iterable |
21 | 22 | from datetime import datetime, timedelta |
22 | | -from typing import TYPE_CHECKING, Any |
| 23 | +from typing import TYPE_CHECKING, Any, ClassVar |
23 | 24 |
|
24 | 25 | from airflow._shared.timezones.timezone import make_aware, parse_timezone |
25 | 26 | from airflow.partition_mappers.base import PartitionMapper |
26 | 27 |
|
27 | 28 | if TYPE_CHECKING: |
28 | 29 | from pendulum import FixedTimezone, Timezone |
29 | 30 |
|
| 31 | + from airflow.partition_mappers.window import Window |
| 32 | + |
30 | 33 |
|
31 | 34 | _STRPTIME_PATTERNS: dict[str, str] = { |
32 | 35 | "%Y": r"\d{4}", |
@@ -419,3 +422,132 @@ def normalize(self, dt: datetime) -> datetime: |
419 | 422 | second=0, |
420 | 423 | microsecond=0, |
421 | 424 | ) |
| 425 | + |
| 426 | + |
| 427 | +class FanOutMapper(PartitionMapper): |
| 428 | + """ |
| 429 | + Partition mapper that fans one upstream key out into multiple downstream keys. |
| 430 | +
|
| 431 | + Compose an ``upstream_mapper`` (parses the coarse upstream key and |
| 432 | + normalizes it to its period start) with a ``window`` (enumerates the |
| 433 | + members of that period). ``downstream_mapper`` formats each member into a |
| 434 | + downstream key string; if omitted, a default is chosen from the window |
| 435 | + class. |
| 436 | +
|
| 437 | + ``downstream_mapper`` must be passed explicitly for any window without an |
| 438 | + entry in the default table — currently ``HourWindow`` and any custom |
| 439 | + ``Window`` subclass. Constructing a ``FanOutMapper`` for those windows |
| 440 | + without a ``downstream_mapper`` raises ``ValueError`` at Dag-parse time. |
| 441 | +
|
| 442 | + Symmetric to :class:`~airflow.partition_mappers.base.RollupMapper`: rollup |
| 443 | + is N→1 (downstream waits until all members arrive), fan-out is 1→N (one |
| 444 | + upstream event creates one downstream Dag run per member). |
| 445 | +
|
| 446 | + .. code-block:: python |
| 447 | +
|
| 448 | + # Weekly upstream → 7 daily downstream Dag runs |
| 449 | + FanOutMapper(upstream_mapper=StartOfWeekMapper(), window=WeekWindow()) |
| 450 | + """ |
| 451 | + |
| 452 | + # Keep ``FanOutMapper.default_downstream_mapper_by_window_name`` in sync with |
| 453 | + # the SDK copy in |
| 454 | + # ``task-sdk/src/airflow/sdk/definitions/partition_mappers/temporal.py`` — |
| 455 | + # the SDK and core class hierarchies are independent (the SDK cannot import |
| 456 | + # core), so both sides carry the same defaults and the lookup is by class |
| 457 | + # name. When adding a new ``Window`` subclass, extend the table on both |
| 458 | + # sides; a missing entry raises ``ValueError`` at ``FanOutMapper.__init__`` |
| 459 | + # (see ``FanOutMapper._resolve_default_downstream_mapper``). The |
| 460 | + # ``check-partition-mapper-defaults-in-sync`` prek hook enforces that the |
| 461 | + # two tables stay identical. |
| 462 | + default_downstream_mapper_by_window_name: ClassVar[dict[str, type[_BaseTemporalMapper]]] = { |
| 463 | + "DayWindow": StartOfHourMapper, |
| 464 | + "WeekWindow": StartOfDayMapper, |
| 465 | + "MonthWindow": StartOfDayMapper, |
| 466 | + "QuarterWindow": StartOfMonthMapper, |
| 467 | + "YearWindow": StartOfMonthMapper, |
| 468 | + } |
| 469 | + |
| 470 | + @classmethod |
| 471 | + def _resolve_default_downstream_mapper(cls, window: Window) -> PartitionMapper: |
| 472 | + """ |
| 473 | + Return the conventional downstream mapper for *window*. |
| 474 | +
|
| 475 | + Looked up by the window's class **name** rather than identity so that |
| 476 | + the SDK ``Window`` classes (used in Dag-author code) and the core |
| 477 | + ``Window`` classes (used after deserialization) both resolve to the |
| 478 | + same default. Subclasses can extend or override the defaults by |
| 479 | + setting :attr:`default_downstream_mapper_by_window_name` on the |
| 480 | + subclass. |
| 481 | + """ |
| 482 | + mapper_cls = cls.default_downstream_mapper_by_window_name.get(type(window).__name__) |
| 483 | + if mapper_cls is None: |
| 484 | + raise ValueError( |
| 485 | + f"{cls.__name__} has no default downstream_mapper for window type " |
| 486 | + f"{type(window).__name__}; pass downstream_mapper explicitly." |
| 487 | + ) |
| 488 | + return mapper_cls() |
| 489 | + |
| 490 | + def __init__( |
| 491 | + self, |
| 492 | + *, |
| 493 | + upstream_mapper: PartitionMapper, |
| 494 | + window: Window, |
| 495 | + downstream_mapper: PartitionMapper | None = None, |
| 496 | + ) -> None: |
| 497 | + self.upstream_mapper = upstream_mapper |
| 498 | + self.window = window |
| 499 | + self.downstream_mapper = downstream_mapper or self._resolve_default_downstream_mapper(window) |
| 500 | + |
| 501 | + def to_downstream(self, key: str) -> Iterable[str]: |
| 502 | + # Round-trip the upstream key through its mapper to obtain the |
| 503 | + # period-start datetime (decoded form). This keeps the upstream_mapper |
| 504 | + # opaque — we don't need to know whether it's temporal or segment. |
| 505 | + formatted = self.upstream_mapper.to_downstream(key) |
| 506 | + if not isinstance(formatted, str): |
| 507 | + raise TypeError( |
| 508 | + "FanOutMapper.upstream_mapper must produce a single key from " |
| 509 | + "to_downstream; chained fan-out (mapper that itself returns multiple keys) " |
| 510 | + "is not supported." |
| 511 | + ) |
| 512 | + coarse = self.upstream_mapper.decode_downstream(formatted) |
| 513 | + return [_format_with(self.downstream_mapper, item) for item in self.window.to_upstream(coarse)] |
| 514 | + |
| 515 | + def serialize(self) -> dict[str, Any]: |
| 516 | + from airflow.serialization.encoders import encode_partition_mapper, encode_window |
| 517 | + |
| 518 | + return { |
| 519 | + "upstream_mapper": encode_partition_mapper(self.upstream_mapper), |
| 520 | + "window": encode_window(self.window), |
| 521 | + "downstream_mapper": encode_partition_mapper(self.downstream_mapper), |
| 522 | + } |
| 523 | + |
| 524 | + @classmethod |
| 525 | + def deserialize(cls, data: dict[str, Any]) -> PartitionMapper: |
| 526 | + from airflow.serialization.decoders import decode_partition_mapper, decode_window |
| 527 | + |
| 528 | + return cls( |
| 529 | + upstream_mapper=decode_partition_mapper(data["upstream_mapper"]), |
| 530 | + window=decode_window(data["window"]), |
| 531 | + downstream_mapper=decode_partition_mapper(data["downstream_mapper"]), |
| 532 | + ) |
| 533 | + |
| 534 | + |
| 535 | +def _format_with(mapper: PartitionMapper, decoded: Any) -> str: |
| 536 | + """ |
| 537 | + Format *decoded* using *mapper*'s downstream format. |
| 538 | +
|
| 539 | + Three-layer fallback, in order: |
| 540 | +
|
| 541 | + 1. *mapper* exposes a callable ``format`` attribute (all temporal mappers do) |
| 542 | + — delegates to ``mapper.format(decoded)``, which uses ``output_format``. |
| 543 | + 2. *decoded* is a ``datetime`` instance — returns ``decoded.isoformat()``, |
| 544 | + producing a stable ISO-8601 string that round-trips via |
| 545 | + ``datetime.fromisoformat``. |
| 546 | + 3. Anything else — ``str(decoded)`` as a last-resort fallback. |
| 547 | + """ |
| 548 | + formatter = getattr(mapper, "format", None) |
| 549 | + if callable(formatter): |
| 550 | + return formatter(decoded) |
| 551 | + if isinstance(decoded, datetime): |
| 552 | + return decoded.isoformat() |
| 553 | + return str(decoded) |
0 commit comments