dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4093508270
##########
airflow-core/src/airflow/serialization/definitions/mappedoperator.py:
##########
@@ -526,6 +587,25 @@ def _(task: SerializedBaseOperator | TaskSDKBaseOperator,
run_id: str, *, sessio
def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *,
session: Session) -> int:
from airflow.serialization.serialized_objects import BaseSerialization,
_ExpandInputRef
+ def _get_parent_count() -> int:
+ if (group := task.get_closest_mapped_task_group()) is None:
+ return 1
+ return get_mapped_ti_count(group, run_id, session=session)
+
+ # See get_parse_time_mapped_ti_count: a batched task's count is fixed by
batch_size alone.
+ if isinstance(task, SerializedMappedOperator):
+ batch_size = task.resolve_batch_size(run_id, session=session)
+ elif isinstance(task.batch_size, int):
+ batch_size = task.batch_size
+ else:
+ # A runtime batch size lives in task_map and is only resolvable
through the serialized
+ # operator; SDK objects only reach here from tests that skip
serialization.
+ raise TypeError(
+ f"runtime batch size of {task.task_id!r} can only be resolved on a
serialized operator"
Review Comment:
Fixed in f55fb22809. Confirmed that `dag.test()` serialises the DAG before
scheduling, so it never reached this branch, but it no longer matters: the
unserialized case now counts the same way. The `task_map` lookup moved into a
shared helper, and the SDK branch reads the live `XComArg`'s reference where
the serialized one dereferences its `_XComRef`. The `task_map` count test now
runs both serialized and unserialized and pins the same outcomes for both.
---
Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting
--
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]
