rjgoyln commented on code in PR #70620:
URL: https://github.com/apache/airflow/pull/70620#discussion_r3674102590
##########
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()
Review Comment:
Benchmarked it — `DISTINCT` should stay. `set()` deduplicates after every
row has crossed the connection and been materialized in Python; `DISTINCT`
does it in the database before the transfer. The duplication factor here is
the number of map indexes of the mapped upstream, so it is exactly the
workload this PR targets.
PostgreSQL 16, probing 3 mapped upstreams:
| map indexes | rows scanned | DISTINCT | plain |
|------------:|-------------:|---------:|--------:|
| 1 | 3 | 0.451ms | 0.441ms |
| 100 | 300 | 0.517ms | 0.778ms |
| 1000 | 3,000 | 1.514ms | 6.009ms |
| 5000 | 15,000 | 4.508ms | 22.593ms |
SQLite shows the same shape (6.2x at 15k rows). Break-even is in the
single-digit-rows range, where the difference is ~0.01ms — noise. Above
that `DISTINCT` wins by 4-6x.
---
Drafted-by: Claude Code (Opus 5)
--
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]
