uranusjr commented on code in PR #70620:
URL: https://github.com/apache/airflow/pull/70620#discussion_r3671835852
##########
airflow-core/src/airflow/serialization/definitions/xcom_arg.py:
##########
@@ -145,14 +154,102 @@ def iter_references(self) -> Iterator[tuple[Operator,
str]]:
yield from arg.iter_references()
+def _match_referenced_tasks(model: Any, keys: Iterable[tuple[str, str]]) ->
tuple[Any, Any]:
+ """
+ Build filters selecting rows of ``model`` for any of the given ``(dag_id,
task_id)``.
+
+ Matching the two columns independently rather than as a row value keeps
the query
+ portable. That is wider than the key set only if the keys span several
Dags, which
+ the callers never do, and callers read results back by exact key
regardless.
+ """
+ dag_ids, task_ids = zip(*keys)
+ return model.dag_id.in_(set(dag_ids)), model.task_id.in_(set(task_ids))
+
+
+def prefetch_map_lengths(
+ xcom_args: Iterable[SchedulerXComArg], run_id: str, *, session: Session
+) -> dict[tuple[str, str], int]:
+ """
+ Resolve the map length of every task referenced by ``xcom_args`` in bulk.
+
+ Passing the result to :func:`get_task_map_length` as ``lengths`` keeps the
number of
+ queries constant no matter how many arguments -- and how many tasks nested
inside
+ ``zip()``/``concat()`` arguments -- have to be resolved.
+
+ Tasks whose length is not known yet are absent from the result, mirroring
the
+ ``None`` that :func:`get_task_map_length` returns for them.
+ """
+ from airflow.models.taskinstance import TaskInstance
+ from airflow.models.taskmap import TaskMap
+ from airflow.models.xcom import XComModel
+ from airflow.serialization.definitions.mappedoperator import is_mapped
+
+ operators = {(op.dag_id, op.task_id): op for arg in xcom_args for op, _ in
arg.iter_references()}
+ if not operators:
+ return {}
+ mapped = {key for key, op in operators.items() if is_mapped(op)}
+ unmapped = operators.keys() - mapped
+
+ lengths: dict[tuple[str, str], int] = {}
+ if unmapped:
+ rows = session.execute(
+ select(TaskMap.dag_id, TaskMap.task_id, TaskMap.length).where(
+ TaskMap.run_id == run_id,
+ TaskMap.map_index < 0,
+ *_match_referenced_tasks(TaskMap, unmapped),
+ )
+ )
+ lengths.update({(dag_id, task_id): length for dag_id, task_id, length
in rows})
+ if mapped:
+ unfinished = set(
+ session.execute(
+ select(TaskInstance.dag_id, TaskInstance.task_id)
+ .where(
+ TaskInstance.run_id == run_id,
+ *_match_referenced_tasks(TaskInstance, mapped),
+ # Special NULL treatment is needed because 'state' can be
NULL.
+ # The "IN" part would produce "NULL NOT IN ..." and
eventually
+ # "NULl = NULL", which is a big no-no in SQL.
+ or_(
+ TaskInstance.state.is_(None),
+ TaskInstance.state.in_(s.value for s in
State.unfinished if s is not None),
+ ),
+ )
+ .distinct()
+ )
+ )
+ if finished := mapped - unfinished:
+ counts = {
+ (dag_id, task_id): count
+ for dag_id, task_id, count in session.execute(
+ select(XComModel.dag_id, XComModel.task_id,
func.count(XComModel.map_index))
+ .where(
+ XComModel.run_id == run_id,
+ XComModel.map_index >= 0,
+ XComModel.key == XCOM_RETURN_KEY,
+ *_match_referenced_tasks(XComModel, finished),
+ )
+ .group_by(XComModel.dag_id, XComModel.task_id)
+ )
+ }
+ # A finished mapped task that pushed nothing has no row to group,
but its
+ # length is a known zero rather than an unresolved value.
+ lengths.update({key: counts.get(key, 0) for key in finished})
+ return lengths
+
+
@singledispatch
-def get_task_map_length(xcom_arg: SchedulerXComArg, run_id: str, *, session:
Session) -> int | None:
+def get_task_map_length(
+ xcom_arg: SchedulerXComArg, run_id: str, *, session: Session, lengths:
MapLengths | None = None
Review Comment:
Please always put `session` last.
--
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]
