Rewrite decorated task mapping

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.

[uranusjr]

Need to get #21210 merged first and rebase this. Done

This rewrites how map() is implemented on _TaskDecorator. The previous implementation incorrectly assumed mapping a task decorator works like mapping a traditional operator, but in fact they have entirely different semantics.

When mapping a traditional operator, FooOperator.map(my_arg=[1, 2, 3]), the argument on FooOperator is mapped. But when mapping a task decorator, my_task.map(my_val=[1, 2, 3]), it’s the argument on the function wrapped by the task object being mapped. Therefore, mapped (and also partial-ed) values for my_val should go into the DecoratedOperator’s op_kwargs (and op_args) arguments instead.

An end-to-end test is provided to validate the simplest case of this. Also fixed a test where the end-to-end test becomes broken after running a DAG (because I need to run a second DAG!)

[ashb]

The serialization needs some work for this approach. Here is an incomplete test:

deftest_mapped_decorator_serde():
fromairflow.models.xcom_argimportXComArgfromairflow.decoratorsimporttaskwithDAG("test-dag", start_date=datetime(2020, 1, 1)) asdag:
task1=BaseOperator(task_id="op1")
xcomarg=XComArg(task1, "test_key")
@task(retry_delay=30)defx(arg1, arg2):
...
real_op=x.partial(arg1=1).map(arg2=xcomarg).operatorserialized=SerializedBaseOperator._serialize(real_op)
assertserialized== {
'_is_dummy': False,
'_is_mapped': True,
'_task_module': 'airflow.decorators.python',
'_task_type': '_PythonDecoratedOperator',
'downstream_task_ids': [],
'partial_kwargs': {
'multiple_outputs': False,
'op_args': [],
'op_kwargs': {
'arg1': [
1,
2,
{"__type": "dict", "__var": {'a': 'b'}},
],
},
'retry_delay': 30,
},
'mapped_kwargs': {
'op_args': [],
# We don't need the __type/__var here!'op_kwargs': {
'arg2': {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'test_key'}},
}
},
'task_id': 'x',
'template_ext': [],
'template_fields': ['op_args', 'op_kwargs'],
# We don't want to include the python source code in the serialized representation# TODO? Where does `retry_delay` go? We need to separate 
}
op=SerializedBaseOperator.deserialize_operator(serialized)
assertisinstance(op, MappedOperator)
assertop.depsisMappedOperator.DEFAULT_DEPS# TODO: add some more asserts here

The serialization

Trying to call real_op.unmap() also throws a key error:

tests/serialization/test_dag_serialization.py:1670: in test_mapped_decorator_serde
real_op.unmap()
airflow/models/baseoperator.py:1882: in unmap
dag._remove_task(self.task_id)
airflow/models/dag.py:2167: in _remove_task
task = self.task_dict.pop(task_id)
E KeyError: 'x'

[uranusjr]

Serialisation format is revised. The KeyError is also fixed by implementing slightly smarter logic to merge op_kwargs from mapped_kwargs and partial_kwargs.

[uranusjr]

We should also implement some validation when a DAG is parsed to make sure the user pass reasonable values to partial_kwargs. This will be done in a separate PR.

[ashb]

Suggested change

	assertdeserialized.retry_delay==timedelta(seconds=30)

(give or take.)

[uranusjr]

At this point retry_delay is still 30 verbatim; it become a timedelta only after unmapped. I’ll add a test elsewhere for this.

[ashb]

It should be deserialized as a timedelta -- it's fine for this value since the constructor for BaseOperator handles it, but other variables might behave differently

[ashb]

start_date for instance.

[uranusjr]

I added a test in tests/decorators/test_python.py for this

[uranusjr]

start_date for instance.

Hm some refactoring would be called for to extract the logic out of BaseOperator for reuse. (Note that this affects non-decorator MappedOperator as well.) I think this should be done in a separate PR.

[ashb]

Yeah agreed. We can fix this up later, this unblocks a lot.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[uranusjr]

Static check failures fixed in #21480.