Lee-W commented on code in PR #64571:
URL: https://github.com/apache/airflow/pull/64571#discussion_r3116753319
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/ui/assets.py:
##########
@@ -115,6 +119,76 @@ def next_run_assets(
if not event.pop("queued", None):
event["lastUpdate"] = None
+ # For partitioned DAGs: enrich events with per-asset received/required
counts,
+ # using to_upstream for rollup mappers, and fix lastUpdate for partial
receipt.
+ if is_partitioned:
+ timetable = _load_timetable(dag_id, session)
+
+ pending_apdr = session.execute(
+ select(AssetPartitionDagRun.id, AssetPartitionDagRun.partition_key)
+ .where(
+ AssetPartitionDagRun.target_dag_id == dag_id,
+ AssetPartitionDagRun.created_dag_run_id.is_(None),
+ )
+ .order_by(AssetPartitionDagRun.created_at.desc())
+ .limit(1)
+ ).one_or_none()
+
+ if pending_apdr is not None:
+ # Collect received upstream partition keys per asset for this
partition run.
+ # Use a set to deduplicate: multiple events for the same key count
as one.
+ received_keys_by_asset: dict[int, set[str]] = {}
+ for row in session.execute(
+ select(
+ PartitionedAssetKeyLog.asset_id,
+ PartitionedAssetKeyLog.source_partition_key,
+ ).where(PartitionedAssetKeyLog.asset_partition_dag_run_id ==
pending_apdr.id)
+ ):
+ received_keys_by_asset.setdefault(row.asset_id,
set()).add(row.source_partition_key or "")
+
+ for event in events:
+ asset_id = event["id"]
+ received_keys = sorted(received_keys_by_asset.get(asset_id,
set()))
+ required_keys: list[str] = [pending_apdr.partition_key]
+ is_rollup = False
+ if timetable is not None:
+ with suppress(Exception):
+ mapper = timetable.get_partition_mapper(
+ name=event["name"],
+ uri=event["uri"],
+ )
+ if isinstance(mapper, RollupMapper):
+ required_keys =
sorted(mapper.to_upstream(pending_apdr.partition_key))
+ is_rollup = True
+ received_count = len(received_keys)
+ required_count = len(required_keys)
+ event["receivedCount"] = received_count
+ event["requiredCount"] = required_count
+ event["receivedKeys"] = received_keys
+ event["requiredKeys"] = required_keys
+ event["isRollup"] = is_rollup
+ # Only show lastUpdate when all required upstream keys are
received
+ if received_count < required_count:
+ event["lastUpdate"] = None
Review Comment:
creating a data model for it. the camelCase seems to be something messed up
during squashing and fixup... sorry for the mess
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/ui/assets.py:
##########
@@ -115,6 +119,76 @@ def next_run_assets(
if not event.pop("queued", None):
event["lastUpdate"] = None
+ # For partitioned DAGs: enrich events with per-asset received/required
counts,
+ # using to_upstream for rollup mappers, and fix lastUpdate for partial
receipt.
+ if is_partitioned:
+ timetable = _load_timetable(dag_id, session)
+
+ pending_apdr = session.execute(
+ select(AssetPartitionDagRun.id, AssetPartitionDagRun.partition_key)
+ .where(
+ AssetPartitionDagRun.target_dag_id == dag_id,
+ AssetPartitionDagRun.created_dag_run_id.is_(None),
+ )
+ .order_by(AssetPartitionDagRun.created_at.desc())
+ .limit(1)
+ ).one_or_none()
+
+ if pending_apdr is not None:
+ # Collect received upstream partition keys per asset for this
partition run.
+ # Use a set to deduplicate: multiple events for the same key count
as one.
+ received_keys_by_asset: dict[int, set[str]] = {}
+ for row in session.execute(
+ select(
+ PartitionedAssetKeyLog.asset_id,
+ PartitionedAssetKeyLog.source_partition_key,
+ ).where(PartitionedAssetKeyLog.asset_partition_dag_run_id ==
pending_apdr.id)
+ ):
+ received_keys_by_asset.setdefault(row.asset_id,
set()).add(row.source_partition_key or "")
+
+ for event in events:
+ asset_id = event["id"]
+ received_keys = sorted(received_keys_by_asset.get(asset_id,
set()))
+ required_keys: list[str] = [pending_apdr.partition_key]
+ is_rollup = False
+ if timetable is not None:
+ with suppress(Exception):
+ mapper = timetable.get_partition_mapper(
+ name=event["name"],
+ uri=event["uri"],
+ )
+ if isinstance(mapper, RollupMapper):
+ required_keys =
sorted(mapper.to_upstream(pending_apdr.partition_key))
Review Comment:
replaced it with `is_rollup`. missed a few
##########
airflow-core/src/airflow/example_dags/example_asset_partition.py:
##########
@@ -225,3 +227,68 @@ def regional_stats_breakdown():
keys belong to a fixed set of allowed values (``us``, ``eu``, ``apac``)
rather than time-based partitions.
"""
pass
+
+
+daily_sales = Asset(uri="s3://sales/daily", name="daily_sales")
+daily_costs = Asset(uri="s3://costs/daily", name="daily_costs")
+
+
+with DAG(
+ dag_id="produce_daily_sales",
+ schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"),
+ tags=["sales", "ingestion"],
+):
+ """Produce daily sales data partitioned by date."""
+
+ @task(outlets=[daily_sales])
+ def upload_daily_sales(dag_run=None):
Review Comment:
My editor still complains 🤔 but pyright complains on lots of things. will
check whether mypy likes it better
##########
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:
is there any other day-based format in your mind that we should cover? I
kinda feel like there might be false positive. but i can give it a try if
there's more concrete examples. thanks!
--
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]
