Repository: apache/airflow

Issue:
Dynamic Task Mapping not working with op_kwargs in PythonOperator
### Apache Airflow version

2.3.0 (latest released)

### What happened

The following DAG was written and expected to generate 3 tasks (one for each string in the list)

**dag_code**
```python
import logging
from airflow.decorators import dag, task
from airflow.operators.python import PythonOperator

from airflow.utils.dates import datetime


def log_strings_operator(string, *args, **kwargs):
    logging.info("we've made it into the method")
    logging.info(f"operator log - {string}")


@dag(
    dag_id='dynamic_dag_test',
    schedule_interval=None,
    start_date=datetime(2021, 1, 1),
    catchup=False,
    tags=['example', 'dynamic_tasks']
)
def tutorial_taskflow_api_etl():
    op2 = (PythonOperator
           .partial(task_id="logging_with_operator_task",
                    python_callable=log_strings_operator)
           .expand(op_kwargs=[{"string": "a"}, {"string": "b"}, {"string": "c"}]))

    return op2


tutorial_etl_dag = tutorial_taskflow_api_etl()
```

**error message**
```python
Broken DAG: [/usr/local/airflow/dags/dynamic_dag_test.py] Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/airflow/serialization/serialized_objects.py", line 343, in _serialize
    return SerializedBaseOperator.serialize_mapped_operator(var)
  File "/usr/local/lib/python3.9/site-packages/airflow/serialization/serialized_objects.py", line 608, in serialize_mapped_operator
    assert op_kwargs[Encoding.TYPE] == DAT.DICT
TypeError: list indices must be integers or slices, not Encoding

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/airflow/serialization/serialized_objects.py", line 1105, in to_dict
    json_dict = {"__version": cls.SERIALIZER_VERSION, "dag": cls.serialize_dag(var)}
  File "/usr/local/lib/python3.9/site-packages/airflow/serialization/serialized_objects.py", line 1013, in serialize_dag
    raise SerializationError(f'Failed to serialize DAG {dag.dag_id!r}: {e}')
airflow.exceptions.SerializationError: Failed to serialize DAG 'dynamic_dag_test': list indices must be integers or slices, not Encoding
```

### What you think should happen instead

Dag should contain 1 task `logging_with_operator_task` that contains 3 indices

### How to reproduce

copy/paste dag code into a dag file and run on airflow 2.3.0. Airflow UI will flag the error

### Operating System

Debian

### Versions of Apache Airflow Providers

_No response_

### Deployment

Astronomer

### 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

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

location of the error itself: 
https://github.com/apache/airflow/blob/main/airflow/serialization/serialized_objects.py#L590-L620
@erdos2n  i guess, you need to pass `op_kwargs` as a `dict` : 
```
  .expand(op_kwargs={"string":[ "a", "b","c"]}))
```
@snjypl I tried this and it still didn't work, were you able to get it to work? 
yes @erdos2n  i was able to get that work. can you please share the error you are getting? 
@snjypl here is the code and error message when it runs in Airflow. So the UI Error does disappear, but it still isn't working. 

```python
import logging
from airflow.decorators import dag, task
from airflow.operators.python import PythonOperator

from airflow.utils.dates import datetime

def log_strings(string):
    logging.info(f"here is the string = {string}")

@dag(
    dag_id='dynamic_dag_test',
    schedule_interval=None,
    start_date=datetime(2021, 1, 1),
    catchup=False,
    tags=['example', 'dynamic_tasks']
)
def tutorial_taskflow_api_etl():
    op2 = (PythonOperator
           .partial(task_id="logging_with_operator_task",
                    python_callable=log_strings)
           .expand(op_kwargs={"string":["a", "b", "c"]}))

    return op2


tutorial_etl_dag = tutorial_taskflow_api_etl()
```

**Error message:**
```python
[2022-05-22, 17:06:49 UTC] {taskinstance.py:1376} INFO - Executing <Mapped(PythonOperator): logging_with_operator_task> on 2022-05-22 17:06:48.858540+00:00
[2022-05-22, 17:06:49 UTC] {standard_task_runner.py:52} INFO - Started process 290 to run task
[2022-05-22, 17:06:49 UTC] {standard_task_runner.py:79} INFO - Running: ['airflow', 'tasks', 'run', 'dynamic_dag_test', 'logging_with_operator_task', 'manual__2022-05-22T17:06:48.858540+00:00', '--job-id', '41', '--raw', '--subdir', 'DAGS_FOLDER/dynamic_dag_test.py', '--cfg-path', '/tmp/tmp1yeh1bff', '--map-index', '0', '--error-file', '/tmp/tmp2mzffl7i']
[2022-05-22, 17:06:49 UTC] {standard_task_runner.py:80} INFO - Job 41: Subtask logging_with_operator_task
[2022-05-22, 17:06:49 UTC] {task_command.py:369} INFO - Running <TaskInstance: dynamic_dag_test.logging_with_operator_task manual__2022-05-22T17:06:48.858540+00:00 map_index=0 [running]> on host 5b49114612fc
[2022-05-22, 17:06:49 UTC] {taskinstance.py:1568} INFO - Exporting the following env vars:
AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=dynamic_dag_test
AIRFLOW_CTX_TASK_ID=logging_with_operator_task
AIRFLOW_CTX_EXECUTION_DATE=2022-05-22T17:06:48.858540+00:00
AIRFLOW_CTX_TRY_NUMBER=1
AIRFLOW_CTX_DAG_RUN_ID=manual__2022-05-22T17:06:48.858540+00:00
[2022-05-22, 17:06:49 UTC] {taskinstance.py:1888} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/airflow/operators/python.py", line 168, in execute
    context_merge(context, self.op_kwargs, templates_dict=self.templates_dict)
  File "/usr/local/lib/python3.9/site-packages/airflow/utils/context.py", line 256, in context_merge
    context.update(*args, **kwargs)
  File "/usr/local/lib/python3.9/_collections_abc.py", line 946, in update
    for key, value in other:
ValueError: too many values to unpack (expected 2)
[2022-05-22, 17:06:50 UTC] {taskinstance.py:1394} INFO - Marking task as FAILED. dag_id=dynamic_dag_test, task_id=logging_with_operator_task, map_index=0, execution_date=20220522T170648, start_date=20220522T170649, end_date=20220522T170650
[2022-05-22, 17:06:50 UTC] {standard_task_runner.py:92} ERROR - Failed to execute job 41 for task logging_with_operator_task (too many values to unpack (expected 2); 290)
[2022-05-22, 17:06:50 UTC] {local_task_job.py:156} INFO - Task exited with return code 1
[2022-05-22, 17:06:50 UTC] {local_task_job.py:273} INFO - 0 downstream tasks scheduled from follow-on schedule check
```
Your original syntax is correct. (https://github.com/apache/airflow/issues/23833#issuecomment-1133900501 is wrong.) This looks like a bug in the serialisation code.
@uranusjr I agree, the first syntax best aligns with the current implementation of the dynamic mapping methods. 

Resolving patch:
diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -39,7 +39,7 @@
 import attr
 import typing_extensions
 
-from airflow.compat.functools import cache, cached_property
+from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException
 from airflow.models.abstractoperator import DEFAULT_RETRIES, DEFAULT_RETRY_DELAY
 from airflow.models.baseoperator import (
@@ -413,14 +413,6 @@ class DecoratedMappedOperator(MappedOperator):
     def __hash__(self):
         return id(self)
 
-    @classmethod
-    @cache
-    def get_serialized_fields(cls):
-        # The magic super() doesn't work here, so we use the explicit form.
-        # Not using super(..., cls) to work around pyupgrade bug.
-        sup = super(DecoratedMappedOperator, DecoratedMappedOperator)
-        return sup.get_serialized_fields() | {"mapped_op_kwargs"}
-
     def __attrs_post_init__(self):
         # The magic super() doesn't work here, so we use the explicit form.
         # Not using super(..., self) to work around pyupgrade bug.
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -318,6 +318,7 @@ def get_serialized_fields(cls):
             "dag",
             "deps",
             "is_mapped",
+            "mapped_kwargs",  # This is needed to be able to accept XComArg.
             "subdag",
             "task_group",
             "upstream_task_ids",
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -16,7 +16,7 @@
 # under the License.
 
 """Serialized DAG and BaseOperator"""
-import contextlib
+
 import datetime
 import enum
 import logging
@@ -591,6 +591,9 @@ def task_type(self, task_type: str):
     def serialize_mapped_operator(cls, op: MappedOperator) -> Dict[str, Any]:
         serialized_op = cls._serialize_node(op, include_deps=op.deps is MappedOperator.deps_for(BaseOperator))
 
+        # Handle mapped_kwargs and mapped_op_kwargs.
+        serialized_op[op._expansion_kwargs_attr] = cls._serialize(op._get_expansion_kwargs())
+
         # Simplify partial_kwargs by comparing it to the most barebone object.
         # Remove all entries that are simply default values.
         serialized_partial = serialized_op["partial_kwargs"]
@@ -602,20 +605,6 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> Dict[str, Any]:
             if v == default:
                 del serialized_partial[k]
 
-        # Simplify op_kwargs format. It must be a dict, so we flatten it.
-        with contextlib.suppress(KeyError):
-            op_kwargs = serialized_op["mapped_kwargs"]["op_kwargs"]
-            assert op_kwargs[Encoding.TYPE] == DAT.DICT
-            serialized_op["mapped_kwargs"]["op_kwargs"] = op_kwargs[Encoding.VAR]
-        with contextlib.suppress(KeyError):
-            op_kwargs = serialized_op["partial_kwargs"]["op_kwargs"]
-            assert op_kwargs[Encoding.TYPE] == DAT.DICT
-            serialized_op["partial_kwargs"]["op_kwargs"] = op_kwargs[Encoding.VAR]
-        with contextlib.suppress(KeyError):
-            op_kwargs = serialized_op["mapped_op_kwargs"]
-            assert op_kwargs[Encoding.TYPE] == DAT.DICT
-            serialized_op["mapped_op_kwargs"] = op_kwargs[Encoding.VAR]
-
         serialized_op["_is_mapped"] = True
         return serialized_op
 
@@ -751,15 +740,7 @@ def populate_operator(cls, op: Operator, encoded_op: Dict[str, Any]) -> None:
                 v = cls._deserialize_deps(v)
             elif k == "params":
                 v = cls._deserialize_params_dict(v)
-            elif k in ("mapped_kwargs", "partial_kwargs"):
-                if "op_kwargs" not in v:
-                    op_kwargs: Optional[dict] = None
-                else:
-                    op_kwargs = {arg: cls._deserialize(value) for arg, value in v.pop("op_kwargs").items()}
-                v = {arg: cls._deserialize(value) for arg, value in v.items()}
-                if op_kwargs is not None:
-                    v["op_kwargs"] = op_kwargs
-            elif k == "mapped_op_kwargs":
+            elif k == "partial_kwargs":
                 v = {arg: cls._deserialize(value) for arg, value in v.items()}
             elif k in cls._decorated_fields or k not in op.get_serialized_fields():
                 v = cls._deserialize(v)


Test patch:
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -1688,18 +1688,13 @@ def test_mapped_operator_serde():
         '_task_type': 'BashOperator',
         'downstream_task_ids': [],
         'mapped_kwargs': {
-            'bash_command': [
-                1,
-                2,
-                {"__type": "dict", "__var": {'a': 'b'}},
-            ]
+            "__type": "dict",
+            "__var": {'bash_command': [1, 2, {"__type": "dict", "__var": {'a': 'b'}}]},
         },
         'partial_kwargs': {
             'executor_config': {
                 '__type': 'dict',
-                '__var': {
-                    'dict': {"__type": "dict", "__var": {'sub': 'value'}},
-                },
+                '__var': {'dict': {"__type": "dict", "__var": {'sub': 'value'}}},
             },
         },
         'task_id': 'a',
@@ -1744,7 +1739,10 @@ def test_mapped_operator_xcomarg_serde():
         '_task_module': 'tests.test_utils.mock_operators',
         '_task_type': 'MockOperator',
         'downstream_task_ids': [],
-        'mapped_kwargs': {'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'return_value'}}},
+        'mapped_kwargs': {
+            "__type": "dict",
+            "__var": {'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'return_value'}}},
+        },
         'partial_kwargs': {},
         'task_id': 'task_2',
         'template_fields': ['arg1', 'arg2'],
@@ -1825,13 +1823,18 @@ def x(arg1, arg2, arg3):
         'downstream_task_ids': [],
         'partial_kwargs': {
             'op_args': [],
-            'op_kwargs': {'arg1': [1, 2, {"__type": "dict", "__var": {'a': 'b'}}]},
+            'op_kwargs': {
+                '__type': 'dict',
+                '__var': {'arg1': [1, 2, {"__type": "dict", "__var": {'a': 'b'}}]},
+            },
             'retry_delay': {'__type': 'timedelta', '__var': 30.0},
         },
-        'mapped_kwargs': {},
         'mapped_op_kwargs': {
-            'arg2': {"__type": "dict", "__var": {'a': 1, 'b': 2}},
-            'arg3': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'return_value'}},
+            "__type": "dict",
+            "__var": {
+                'arg2': {"__type": "dict", "__var": {'a': 1, 'b': 2}},
+                'arg3': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'return_value'}},
+            },
         },
         'operator_extra_links': [],
         'ui_color': '#ffefeb',