uranusjr commented on code in PR #64571:
URL: https://github.com/apache/airflow/pull/64571#discussion_r3109780671
##########
airflow-core/src/airflow/partition_mappers/temporal.py:
##########
@@ -99,27 +102,135 @@ def normalize(self, dt: datetime) -> datetime:
class StartOfWeekMapper(_BaseTemporalMapper):
- """Map a time-based partition key to week."""
+ """Map a time-based partition key to the start of its week."""
default_output_format = "%Y-%m-%d (W%V)"
+ def __init__(
+ self,
+ *,
+ week_start: int = 0,
+ timezone: str | Timezone | FixedTimezone = "UTC",
+ input_format: str = "%Y-%m-%dT%H:%M:%S",
+ output_format: str | None = None,
+ ) -> None:
+ super().__init__(timezone=timezone, input_format=input_format,
output_format=output_format)
+ self.week_start = week_start # 0 = Monday (ISO default), 6 = Sunday
+
def normalize(self, dt: datetime) -> datetime:
- start = dt - timedelta(days=dt.weekday())
+ days_since_start = (dt.weekday() - self.week_start) % 7
+ start = dt - timedelta(days=days_since_start)
return start.replace(hour=0, minute=0, second=0, microsecond=0)
+ def serialize(self) -> dict[str, Any]:
+ return {**super().serialize(), "week_start": self.week_start}
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> PartitionMapper:
+ return cls(
+ week_start=data.get("week_start", 0),
+ timezone=parse_timezone(data.get("timezone", "UTC")),
+ input_format=data["input_format"],
+ output_format=data["output_format"],
+ )
+
+
+class WeeklyRollupMapper(StartOfWeekMapper, RollupMapper):
+ """
+ Map a time-based partition key to the start of its week, requiring all 7
daily keys.
+
+ Use this when a partitioned Dag should only run once every daily asset
partition
+ for a full week has been produced. Configure ``week_start`` to set which
day begins
+ the week (0 = Monday, 6 = Sunday).
+ """
+
+ def __init__(self, **kwargs) -> None:
+ super().__init__(**kwargs)
+ if "%Y-%m-%d" not in self.output_format:
+ raise ValueError(
+ f"WeeklyRollupMapper requires output_format to contain
'%Y-%m-%d' so that "
+ f"to_upstream() can recover the week-start date, got:
{self.output_format!r}"
+ )
Review Comment:
I don’t like this very strictly set how the partition key can be. I feel it
should be possible to use some kind of algorithm to recover anything that is
day-based…
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
