Skip to content

Split out and handle 'params' in mapped operator #26100

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
uranusjr merged 2 commits into from
Oct 6, 2022
Merged

Split out and handle 'params' in mapped operator #26100

uranusjr merged 2 commits into from
Oct 6, 2022

Conversation

@uranusjr
Copy link
Member

@uranusjr uranusjr commented Sep 1, 2022 •
edited
Loading

Allow using params in expanding kwargs. Fix #24014.

Syntax unlocked:

# Two tasks with params={"a": 1} and params={"b": 2}.
MyOperator.partial(task_id="t").expand(params=[{"a": 1}, {"b": 2}])

# Two tasks with params={"a": 0, "b": 1} and params={"a": 0, "b": 2}.
# DAG-level params are also merged. If there are duplicated keys, the
# priority is mapped -> task-partial -> DAG.
MyOperator.partial(task_id="t", params={"a": 0}).expand(params=[{"b": 1}, {"b": 2}])

@uranusjr uranusjr added this to the Airflow 2.4.0 milestone Sep 1, 2022
@uranusjr uranusjr force-pushed the aip-42-task-param-expand branch 5 times, most recently from 5824f87 to 7b2115b Compare September 7, 2022 21:59
uranusjr
uranusjr commented Sep 7, 2022
View reviewed changes
Copy link
Member Author

@uranusjr uranusjr left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some notes explaining the implementation, hopefully this is easier to review with them.

airflow/models/mappedoperator.py Outdated
Copy link
Member Author

@uranusjr uranusjr Sep 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be able to expand a task against params, we need to be able to “merge” default DAG- and task-level params with the user-mapped params. So params is split out of other partial kwargs and treated specially. Partial kwargs should never contain the "params" key.

airflow/models/baseoperator.py Outdated
Copy link
Member Author

@uranusjr uranusjr Sep 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So we don’t set the default params in partial_kwargs…

airflow/models/baseoperator.py Outdated
Copy link
Member Author

@uranusjr uranusjr Sep 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

… but pass them to this separate attribute.

airflow/decorators/base.py Outdated
Copy link
Member Author

@uranusjr uranusjr Sep 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And this line is deleted since it’s now no-op.

airflow/models/mappedoperator.py Outdated
Comment on lines -568 to 573
Copy link
Member Author

@uranusjr uranusjr Sep 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actual merging happens here. The merged params is put last so it overrides the incomplete mapped_kwargs["params"].

airflow/models/baseoperator.py Outdated
Comment on lines 1186 to 1196
Copy link
Member Author

@uranusjr uranusjr Sep 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This does not change any behaviour, I only did it so the behaviour is easier to explain…

airflow/models/taskinstance.py Outdated
Comment on lines 2184 to 2188
Copy link
Member Author

@uranusjr uranusjr Sep 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

… here. Previously it’s difficult to follow when rendered_task is None and when it’s not; now it’s simple—a mapped operator (MappedOperator subclases) performs unmapping and returns the unmapped task, while a non-mapped operator (BaseOperator subclasses) return None because no unmapping is needed.

@uranusjr uranusjr marked this pull request as ready for review September 7, 2022 22:49
@ashb ashb modified the milestones: Airflow 2.4.0, Airflow 2.5.0 Sep 8, 2022
@uranusjr uranusjr force-pushed the aip-42-task-param-expand branch from 7b2115b to 50c6989 Compare September 14, 2022 09:18
@uranusjr uranusjr force-pushed the aip-42-task-param-expand branch from 50c6989 to d05e7c2 Compare September 22, 2022 04:35
@uranusjr uranusjr mentioned this pull request Sep 26, 2022
Closed
2 tasks
@uranusjr uranusjr marked this pull request as draft September 27, 2022 09:09
@uranusjr
Copy link
Member Author

uranusjr commented Sep 27, 2022 •
edited
Loading

This is going to conflict with #26702 and I want to get that one in first for 2.4.1. Resolved, should be ready for 2.5.0.

@uranusjr
Split out and handle 'params' in mapped operator
ba229ad
@uranusjr
Implement mapping against 'params'
e9b5861 
It turns out we need to update the template context a bit after
unmapping. This also fixes a bug that context["task"] pointed to the
mapped task but context["ti"].task is the unmapped task. They now both
point to the unmapped one.
@uranusjr uranusjr force-pushed the aip-42-task-param-expand branch from d05e7c2 to e9b5861 Compare September 29, 2022 03:42
@uranusjr uranusjr marked this pull request as ready for review September 29, 2022 03:42
@uranusjr uranusjr merged commit 0ec9651 into apache:main Oct 6, 2022
@uranusjr uranusjr deleted the aip-42-task-param-expand branch October 6, 2022 04:47
@ephraimbuddy ephraimbuddy added the type:new-feature Changelog: New Features label Oct 18, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@dstandish dstandish dstandish approved these changes

@kaxil kaxil Awaiting requested review from kaxil

@ashb ashb Awaiting requested review from ashb

@XD-DENG XD-DENG Awaiting requested review from XD-DENG

Assignees

No one assigned

Labels

area:dynamic-task-mapping AIP-42 type:new-feature Changelog: New Features

Projects

None yet

Milestone

Airflow 2.5.0

Development

Successfully merging this pull request may close these issues.

Task mapping against 'params'

4 participants

@uranusjr @dstandish @ashb @ephraimbuddy