Integration between `XComArg` and `@task_group`
- Dominant language
- Python
- Stars
- 46.9k
- Forks
- 17.8k
- Avg merge
- 2d 9h
- Merged PRs (30d)
- 472
Description
### Description
Task groups are allowed to return a `DAGNode`:
https://github.com/apache/airflow/blob/752419dae9d27419019395b1aaf9ea1bf497b5e9/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py#L204-L206
https://github.com/apache/airflow/blob/752419dae9d27419019395b1aaf9ea1bf497b5e9/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py#L55
A typical `DAGNode` is `Operator`. However in taskflow you typically work with `XComArg`, not `Operator`.
```python
from typing import reveal_type
from airflow.sdk import dag, task, task_group
@task_group
def some_group() -> Operator:
@task
def generate_data() -> list[int]:
return [1, 2, 3, 4, 5]
result = generate_data()
reveal_type(result)
```
```
XComArg
```
To pull a `DAGNode` out from taskflow, you could do something like this, but it's digging into the internals:
```python
if isinstance(result, PlainXComArg):
return result.operator
else:
raise ValueError("Expected a PlainXComArg, got something else.")
```
The best solution would be to allow the `@task_group` decorated function to return an `XComArg`, and then the `DAGNode` could be inferred from there. Failing that, just a consistent API for retrieving the operator from an `XComArg` would be ideal.
***
The other issue is at the other end of the task group: when you use it.
```python
@dag(
start_date=datetime(2024, 1, 1),
)
def workflow():
group_result = some_group()
reveal_type(group_result)
```
```
DAGNode
```
This is the opposite problem. Here we have a `DAGNode`, but taskflow wants to deal with `XComArg`. The hacky workaround is to cast it to `Operator` and then use `.output`, but this is the same problem.
So I think the `@task_group` decorated function would ideally return `XComArg` as well.
### Use case/motivation
This would effectively allow you to use the return value of a `@task_group`. Currently task groups work as a kind of subDAG that can't return values.
### Related issues
_No response_
### Are you willing to submit a PR?
- [x] Yes I am willing to submit a PR!
### Code of Conduct
- [x] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)
Contributor guide
Assessment
This issue has not been assessed yet.