Dynamic mapped tasks group arguments are interpreted as MappedArgument when provided to classic operators

[florian-guily]

Apache Airflow version

Other Airflow 2 version (please specify below)

If "Other Airflow 2 version" selected, which one?

2.8.4

What happened?

When using expand_kwargs on a task group, arguments of this task group will not get correctly interpreted when using them in classic operators inside this task group. They will be interpreted as MappedArgument instead of their real values.

What you think should happen instead?

Real value of the mapped task group's argument should be passed to the operators.

How to reproduce

This was originally done with a google cloud operator in a task group, but i managed to reproduce it with a bash operator.

from airflow import DAG
from airflow.decorators import task, task_group
from airflow.operators.bash import BashOperator
from pendulum import datetime


dag = DAG(
     dag_id="airflow_issue_test",
     start_date=datetime(2024, 1, 1, tz='UTC'),
     catchup=False,
     schedule="@daily",
     default_args={"retries": 2},
 )

with dag:

    @task()
    def list_dict_generator():
        my_list = [
                {
                    "project": "my_project",
                    "dataset": f"dataset_{number}",
                    "table": "my_table",
                    "partition_id": "my_partition_id",
                }
                for number in range(10)
            ]
        return my_list
    
    @task_group()
    def my_tg(project, dataset, table_name, partition_id): 
        BashOperator(
            task_id="bash_task",
            bash_command=f"echo {project}.{dataset}.{table_name}${partition_id}",
            env={"MY_VAR": "Hello World"}
        )
    
    partitions_to_delete = list_dict_generator()

    my_tg.expand_kwargs(partitions_to_delete)

Here are the associated logs. You can clearly see that multiple MappedArguments are passed in the echo command, which should not happen.
dag_id=airflow_issue_test_run_id=scheduled__2024-04-14T00_00_00+00_00_task_id=bq_to_gcs_tg.bash_task_map_index=0_attempt=2.log

Operating System

Mac OS Sonoma 14.2.1 (23C71)

Versions of Apache Airflow Providers

apache-airflow-providers-google==10.16.0
apache-airflow-providers-common-sql>=1.11.0

Deployment

Docker-Compose

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg[bot]]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[RNHTTR]

@florian-guily Would you like to be assigned to this issue?

[florian-guily]

@florian-guily Would you like to be assigned to this issue?

I'd like to yes, i'll find the time to resolve it !

[ghost]

Hello there, I've just encountered the same issue.
It doesn't only happen with classic operators, but anything in a task group scope, such as:

@task_group
def process_model_requests(model_name):
    config_results = prepare_pod_config(task_id=f"config_for_{model_name}")
    result_paths = prepare_pods(config_results, f"for_{model_name}")
    config_results >> result_paths

I have the version 2.6.3 though.

[florian-guily]

Hello there, I've just encountered the same issue.
It doesn't only happen with classic operators, but anything in a task group scope, such as:

@task_group
def process_model_requests(model_name):
    config_results = prepare_pod_config(task_id=f"config_for_{model_name}")
    result_paths = prepare_pods(config_results, f"for_{model_name}")
    config_results >> result_paths

I have the version 2.6.3 though.

Interesting, not anything in a task group though, as python operator created with taskflow api seems to work well. Haven't tried with classic python operator.

[ghost]

Interesting, not anything in a task group though, as python operator created with taskflow api seems to work well. Haven't tried with classic python operator.

Yes, sorry. I meant while using dynamic task mapping. Like the original post.

prefix = "process"
task_params = [{"task_id": "_".join([prefix, str(i + 1)]), "model_name": f"model_{i}"} for i in range(3)]
process_model_requests.expand_kwargs(task_params)

[joshtree41]

Noticing this issue as well, tested a task group with task flow API and the arguments get picked up properly, but when passed to traditional operators they don't get picked up properly.

[Locustv2]

More examples:

1. [google-cloud-transfers] Object of type MappedArgument is not JSON serializable #31481
2. Using XComArg in BigQueryInsertJobOperator `params` #39927

[Mrli2016]

Following your method, I found that this problem can be solved using task.bash. example:

from airflow import DAG
from airflow.decorators import task, task_group
from airflow.operators.bash import BashOperator
from pendulum import datetime


dag = DAG(
     dag_id="airflow_issue_test",
     start_date=datetime(2024, 1, 1, tz='UTC'),
     catchup=False,
     schedule="@daily",
     default_args={"retries": 2},
 )

with dag:

    @task()
    def list_dict_generator():
        my_list = [
                {
                    "project": "my_project",
                    "dataset": f"dataset_{number}",
                    "table": "my_table",
                    "partition_id": "my_partition_id",
                }
                for number in range(10)
            ]
        return my_list

   @task.bash
    def bash_task(project, dataset, table_name, partition_id) -> str:
        return f"echo {project}.{dataset}.{table_name}${partition_id}",
    
    @task_group()
    def my_tg(project, dataset, table_name, partition_id): 
        bash_task(project, dataset, table_name, partition_id)
    
    partitions_to_delete = list_dict_generator()

    my_tg.expand_kwargs(partitions_to_delete)

I'm not sure if this code will run, but I made some simple modifications based on the code from the issue poster. However, the idea is like this, and it works on my end.

[pykenny]

Unfortunately @task.bash is only available from Airflow 2.9, while the original post is on 2.8 environment.

[d-callan]

wondering if there is a known work-around in the mean time? or an update on this? @florian-guily ?

[florian-guily]

didn't find a proper workaround to use non-taskflow operators so i just used python operator for everything.

I'm investigating on the fix but i'm fairly new to open source contribution so i'm a bit slow, i need to find time for that also.

[Locustv2]

The only workaround would be to use taskflow or not use mapped argument IMO

[potiuk]

Yes. And you can use simply Python subprocess.* calls to run bash script in such taskflow-decorated method, there is no particular need to use BashOperator for that if you can just run bash command yourself.

[potiuk]

Converting it to a discussion - as it is unlikely to be an "airflow issue"

[Abirami29]

Hello,
May I ask if there are any updates on this please? or will there be any fix for this in future? I am working on a scenario where I need to use a custom operator with dynamic task mapping inside a task group and unable to do so because of this.

[gnuille]

+1

[xelibrion]

Not sure converting the issue to a discussion was warranted - I believe this is a legitimate issue.
I am trying to use a mix of TaskFlow and classic operator tasks using a suggestion from #29202 (comment), and currently is is failing due to the MappedArgument issue.

[xelibrion]

Found an explanation here https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/dynamic-task-mapping.html#value-references-in-a-task-group-function on why this happens

[ghost]

Totally agree, there is no reason on why it should be encapsulated as a MappedArgument object when it's not a XCom or else.

[pykenny]

XCom is designed to pass message between tasks, not between task groups nor across tasks and task groups. So based on this design decision, data serialization/deserialization should occur within a task's scope instead of a task group's. As the documentation mentioned above says, within DAG definition an operator's XCom output is just a "reference", whose actual value is not determined until downstream task completes pulling its upstream XCom.

Rather than turning the original operator into a decorated Python operator and call subprocess within implementation, there are two other workarounds I can think of:

1. Use templates so that the value is evaluated by rendering the template (not intuitive because you have to hard code what to be pulled from XCom, though):

   bash_template = f"""echo \\
   {{ ti.xcom_pull(task_ids='list_dict_generator',map_indexes=ti.map_index)['project'] }}.\\
   {{ ti.xcom_pull(task_ids='list_dict_generator',map_indexes=ti.map_index)['dataset'] }}.\\
   {{ ti.xcom_pull(task_ids='list_dict_generator',map_indexes=ti.map_index)['table_name'] }}$\\
   {{ ti.xcom_pull(task_ids='list_dict_generator',map_indexes=ti.map_index)['partition_id'] }}"""
   
   BashOperator(
       task_id="bash_task",
       bash_command=bash_template,
       env={"MY_VAR": "Hello World"}
   )

2. Adding an upstream python task to formulate the command, then pass the whole output to operator parameter:

   @task(task_id="emit_command")
   def emit_command(project, dataset, table_name, partition_id) -> str:
       return f"echo {project}.{dataset}.{table_name}${partition_id}"
   
   BashOperator(
       task_id="bash_task",
       bash_command=emit_command(project, dataset, table_name, partition_id),
       env={"MY_VAR": "Hello World"}
   )

[pykenny]

Within the decorated task group context, parameter set (project, dataset, table_name, partition_id) seems to be some native Python variables (strings as for the case proposed) that accepts the values unpacked from upstream XCom, but actually they are special objects that instruct tasks defined within the context what should be pulled from XCom (from which upstream task, and value corresponds to which key) and send to the (templated) parameters for running the task. Concatenating them through f-string without task context merely stringifies these objects together so that's why the weird outcome happened.

[pykenny]

Earlier I shared my thoughts on the reply thread, and I think this kind of confusion on writing decorated task group often happens because unlike tasks, task groups are just a set of layered prefixes attached to task ID mainly used for filtering task subset out from a DAG, instead of a task-like entity having relation with all its underlying task entities. And that's also why XComs/TaskFlow only work on task level, but not for task groups.

Making the proposed case work implies redesigning DAG model is needed, in order to make XCom communication possible on task groups.