uranusjr commented on code in PR #70620:
URL: https://github.com/apache/airflow/pull/70620#discussion_r3671822910
##########
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})
Review Comment:
```suggestion
lengths.update((dag_id, task_id): length for dag_id, task_id, length
in rows)
```
--
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]
