You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "kacpermuda (via GitHub)" <gi...@apache.org> on 2026/02/16 14:12:21 UTC
[PR] feat: Add Dag Edge Labels to OpenLineage event [airflow]
kacpermuda opened a new pull request, #62018:
URL: https://github.com/apache/airflow/pull/62018
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
Airflow supports [edge labeling](https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/dags.html#edge-labels), this PR adds this information to OpenLineage event. Each DAG start event with all tasks listed will now also contain edges coming out from these tasks.
##### Was generative AI tooling used to co-author this PR?
No
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Add Dag Edge Labels to OpenLineage event [airflow]
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski merged PR #62018:
URL: https://github.com/apache/airflow/pull/62018
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Add Dag Edge Labels to OpenLineage event [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #62018:
URL: https://github.com/apache/airflow/pull/62018#discussion_r2817717783
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1361,9 +1361,82 @@ def get_dag_job_dependency_facet(
return {}
+def _build_labeled_edge_map(dag: DAG | SerializedDAG) -> dict[str, tuple[dict, dict]]:
+ """
+ Build a mapping of classified labeled edges for every task and group in the DAG.
+
+ Translates Airflow's internal ``dag.edge_info`` (which uses virtual
+ ``downstream_join_id`` / ``upstream_join_id`` keys) into a clean mapping
+ keyed by real ``task_id`` or ``group_id``. Each value is a tuple of
+ ``(downstream_task_edges, downstream_group_edges)`` sorted dicts.
+
+ ``downstream_task_edges``
+ Target is a **specific task**. Keyed by ``task_id``, regardless of
+ whether source and target are in the same group, different groups,
+ or at root level.
+
+ ``downstream_group_edges``
+ Target is a **TaskGroup as a whole**. Keyed by ``group_id``.
+ Only produced by ``task >> TaskGroup`` or ``TaskGroup >> TaskGroup``.
+
+ **How to read the result (for OL event consumers)**
+
+ Look up ``edge_map.get(task_id)`` or ``edge_map.get(group_id)``.
+ Missing keys mean no labeled edges from that source. Unlabeled edges
+ are not included — use ``downstream_task_ids`` for the full dependency
+ graph.
+
+ Key points:
+
+ - Cross-group task-to-task labels (e.g. ``step_2 >> Label >> final``
+ where the tasks are in different groups) appear under the **source
+ group's** key, not the originating task's key, because Airflow promotes
+ labels to the group boundary.
+ - ``task >> TaskGroup`` produces entries in **both** dicts: per-root-task
+ entries in ``downstream_task_edges`` and a group-level entry in
+ ``downstream_group_edges``, all with the same label.
+
+ See ``test_get_tasks_details_with_edge_labels`` for a comprehensive example
+ covering all edge patterns (same-group, cross-group, task-to-TaskGroup,
+ TaskGroup-to-TaskGroup, nested groups, etc.).
+ """
+ edge_info = dag.edge_info or {}
+ if not edge_info:
+ return {}
+
+ upstream_join_id_to_group_id: dict[str, str] = {}
+ downstream_join_id_to_group_id: dict[str, str] = {}
+
+ for tg_id, tg in dag.task_group_dict.items():
+ upstream_join_id_to_group_id[tg.upstream_join_id] = tg_id
+ downstream_join_id_to_group_id[tg.downstream_join_id] = tg_id
+
+ result: dict[str, tuple[dict, dict]] = {}
+ for source_id, edges in edge_info.items():
+ # Resolve downstream_join_id sources to their owning group_id
+ group_id = downstream_join_id_to_group_id.get(source_id)
+ logical_key = group_id if group_id is not None else source_id
+
+ task_edges: dict = {}
+ group_edges: dict = {}
+ for target_id, info in sorted(edges.items()):
+ resolved_group = upstream_join_id_to_group_id.get(target_id)
+ if resolved_group is not None: # If target is a group, classify as group edge
+ group_edges[resolved_group] = info
+ else:
+ task_edges[target_id] = info
+
+ result[logical_key] = (task_edges, group_edges)
+
+ return result
+
+
def _get_tasks_details(dag: DAG | SerializedDAG) -> dict:
- tasks = {
- single_task.task_id: {
+ edge_map = _build_labeled_edge_map(dag)
Review Comment:
Valid point, adjusted
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1373,23 +1446,26 @@ def _get_tasks_details(dag: DAG | SerializedDAG) -> dict:
"is_setup": single_task.is_setup,
"is_teardown": single_task.is_teardown,
"downstream_task_ids": sorted(single_task.downstream_task_ids),
+ "downstream_task_edges": task_edges,
+ "downstream_group_edges": group_edges,
}
- for single_task in sorted(dag.tasks, key=lambda x: x.task_id)
- }
-
return tasks
def _get_task_groups_details(dag: DAG | SerializedDAG) -> dict:
- return {
- tg_id: {
+ edge_map = _build_labeled_edge_map(dag)
+ result = {}
+ for tg_id, tg in dag.task_group_dict.items():
+ task_edges, group_edges = edge_map.get(tg_id, ({}, {}))
+ result[tg_id] = {
"parent_group": tg.parent_group.group_id,
Review Comment:
Yup, adjusted, 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Add Dag Edge Labels to OpenLineage event [airflow]
Posted by "mobuchowski (via GitHub)" <gi...@apache.org>.
mobuchowski commented on code in PR #62018:
URL: https://github.com/apache/airflow/pull/62018#discussion_r2817596758
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1361,9 +1361,82 @@ def get_dag_job_dependency_facet(
return {}
+def _build_labeled_edge_map(dag: DAG | SerializedDAG) -> dict[str, tuple[dict, dict]]:
+ """
+ Build a mapping of classified labeled edges for every task and group in the DAG.
+
+ Translates Airflow's internal ``dag.edge_info`` (which uses virtual
+ ``downstream_join_id`` / ``upstream_join_id`` keys) into a clean mapping
+ keyed by real ``task_id`` or ``group_id``. Each value is a tuple of
+ ``(downstream_task_edges, downstream_group_edges)`` sorted dicts.
+
+ ``downstream_task_edges``
+ Target is a **specific task**. Keyed by ``task_id``, regardless of
+ whether source and target are in the same group, different groups,
+ or at root level.
+
+ ``downstream_group_edges``
+ Target is a **TaskGroup as a whole**. Keyed by ``group_id``.
+ Only produced by ``task >> TaskGroup`` or ``TaskGroup >> TaskGroup``.
+
+ **How to read the result (for OL event consumers)**
+
+ Look up ``edge_map.get(task_id)`` or ``edge_map.get(group_id)``.
+ Missing keys mean no labeled edges from that source. Unlabeled edges
+ are not included — use ``downstream_task_ids`` for the full dependency
+ graph.
+
+ Key points:
+
+ - Cross-group task-to-task labels (e.g. ``step_2 >> Label >> final``
+ where the tasks are in different groups) appear under the **source
+ group's** key, not the originating task's key, because Airflow promotes
+ labels to the group boundary.
+ - ``task >> TaskGroup`` produces entries in **both** dicts: per-root-task
+ entries in ``downstream_task_edges`` and a group-level entry in
+ ``downstream_group_edges``, all with the same label.
+
+ See ``test_get_tasks_details_with_edge_labels`` for a comprehensive example
+ covering all edge patterns (same-group, cross-group, task-to-TaskGroup,
+ TaskGroup-to-TaskGroup, nested groups, etc.).
+ """
+ edge_info = dag.edge_info or {}
+ if not edge_info:
+ return {}
+
+ upstream_join_id_to_group_id: dict[str, str] = {}
+ downstream_join_id_to_group_id: dict[str, str] = {}
+
+ for tg_id, tg in dag.task_group_dict.items():
+ upstream_join_id_to_group_id[tg.upstream_join_id] = tg_id
+ downstream_join_id_to_group_id[tg.downstream_join_id] = tg_id
+
+ result: dict[str, tuple[dict, dict]] = {}
+ for source_id, edges in edge_info.items():
+ # Resolve downstream_join_id sources to their owning group_id
+ group_id = downstream_join_id_to_group_id.get(source_id)
+ logical_key = group_id if group_id is not None else source_id
+
+ task_edges: dict = {}
+ group_edges: dict = {}
+ for target_id, info in sorted(edges.items()):
+ resolved_group = upstream_join_id_to_group_id.get(target_id)
+ if resolved_group is not None: # If target is a group, classify as group edge
+ group_edges[resolved_group] = info
+ else:
+ task_edges[target_id] = info
+
+ result[logical_key] = (task_edges, group_edges)
+
+ return result
+
+
def _get_tasks_details(dag: DAG | SerializedDAG) -> dict:
- tasks = {
- single_task.task_id: {
+ edge_map = _build_labeled_edge_map(dag)
Review Comment:
We're calling `_build_labeled_edge_map` twice - and it's always called together.
Maybe this should be a param to `_get_tasks_details` and `_get_task_groups_details` and just computed in `get_airflow_job_facet` directly?
Also, we might want add error handling to `get_airflow_job_facet`.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1373,23 +1446,26 @@ def _get_tasks_details(dag: DAG | SerializedDAG) -> dict:
"is_setup": single_task.is_setup,
"is_teardown": single_task.is_teardown,
"downstream_task_ids": sorted(single_task.downstream_task_ids),
+ "downstream_task_edges": task_edges,
+ "downstream_group_edges": group_edges,
}
- for single_task in sorted(dag.tasks, key=lambda x: x.task_id)
- }
-
return tasks
def _get_task_groups_details(dag: DAG | SerializedDAG) -> dict:
- return {
- tg_id: {
+ edge_map = _build_labeled_edge_map(dag)
+ result = {}
+ for tg_id, tg in dag.task_group_dict.items():
+ task_edges, group_edges = edge_map.get(tg_id, ({}, {}))
+ result[tg_id] = {
"parent_group": tg.parent_group.group_id,
Review Comment:
BTW: can `tg.parent_group` be None?
Maybe
```
"parent_group": tg.parent_group.group_id if tg.parent_group else None,
```
as we're doing above:
```
"task_group": single_task.task_group.group_id if single_task.task_group else None,
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org