You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "insomnes (via GitHub)" <gi...@apache.org> on 2025/02/13 18:29:42 UTC
[PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
insomnes opened a new pull request, #46728:
URL: https://github.com/apache/airflow/pull/46728
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
This PR adds the ability to pull XCom when passing multiple `map_indexes` on `RuntimeTaskInstance`.
- Update description in doc string
- Pull XCom for each value from the product of all task ids and map indexes
- Fix `None` case for task id
- Expand original tests
closes: #46665
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975369511
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -1075,44 +1076,86 @@ def test_get_variable_from_context(
assert var_from_context == Variable(key="test_key", value=expected_value)
+ @pytest.mark.parametrize(
+ "map_indexes",
+ [
+ pytest.param(-1, id="no_map_index"),
+ pytest.param(1, id="single_map_index"),
+ pytest.param(42, id="not_set_with_index"),
+ pytest.param([0, 1], id="multiple_map_indexes"),
+ pytest.param((0, 1), id="any_iterable"),
+ pytest.param(None, id="map_indexes_none"),
+ pytest.param(-9999, id="not_set_no_index"),
+ ],
Review Comment:
Oh it does, but there are two ways of doing it (42 and -9999) and it's not very clear from looking at the params alone that this is the behaviour.
Lets not have two ways, and instead of using a maigc int to do it, have one of the param values as NOTSET here (still keeping the `map_indexes_kwarg = {}` behaviour you have)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2669719669
@amoghrajesh thanks for the review! No worries about the delay!
I appreciate your time and I’m glad I could contribute to Task SDK development.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1998061122
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -1075,44 +1076,86 @@ def test_get_variable_from_context(
assert var_from_context == Variable(key="test_key", value=expected_value)
+ @pytest.mark.parametrize(
+ "map_indexes",
+ [
+ pytest.param(-1, id="no_map_index"),
+ pytest.param(1, id="single_map_index"),
+ pytest.param(42, id="not_set_with_index"),
+ pytest.param([0, 1], id="multiple_map_indexes"),
+ pytest.param((0, 1), id="any_iterable"),
+ pytest.param(None, id="map_indexes_none"),
+ pytest.param(-9999, id="not_set_no_index"),
+ ],
Review Comment:
I believe this was fixed, so resolving. Feel free to undo it if I misunderstood the comment.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975366078
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -270,13 +271,19 @@ def xcom_pull(
:param run_id: If provided, only pulls XComs from a DagRun w/a matching run_id.
If *None* (default), the run_id of the calling task is used.
+
Review Comment:
Nit: this might break docs rendering
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2721659341
Thank you @kaxil and @amoghrajesh for the review.
> its taken long to get a review here 🗡️
No problem! I just got a bit lost on the best way to work around the current restrictions of the API, because I don't have a full picture.
I've applied the suggestions as is so I resolved the comments for them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1977563768
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
So, as I get it: I should stick to the original behavior/description and use return with lazy sequence to implement it?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975397441
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
Correct me if I am looking at the wrong API, but the `get_xcom` method here:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/execution_api/routes/xcoms.py#L117
Returns only one XCom result by signature and as I understood through code reading:
```python
# The xcom_query allows no map_index to be passed. This endpoint should always return just a single item,
# so we override that query value
xcom_query = xcom_query.filter(BaseXCom.map_index == map_index)
# We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
# This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
# retrieves the raw serialized value from the database. By not relying on `XCom.get_many` or `XCom.get_one`
# (which automatically deserializes using the backend), we avoid potential
# performance hits from retrieving large data files into the API server.
result = xcom_query.limit(1).first()
```
And `XComResult`, `XComResponse` are single-item models, am I missing something? Will it be multiple inside the JSON? My manual tests showed that with the current API call non-providing is not a working strategy, at least back then. I will recheck.
> not passing a map index as "return all XComs for mapped tis"
This behavior is available at `get_xcom_entries` from public xcom API, and I was targeting it in my initial variant, but after testing I understood that different API is used and I cannot rely on this:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/core_api/routes/public/xcom.py#L109
Please, can you guide me on what kind of behavior is expected from this method? I will try to make it work via `LazySequence` if possible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975411795
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -270,13 +271,19 @@ def xcom_pull(
:param run_id: If provided, only pulls XComs from a DagRun w/a matching run_id.
If *None* (default), the run_id of the calling task is used.
+
Review Comment:
Oh! Didn't know that, thank you. Applied. Maybe this can be enforced via linters?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975397441
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
Correct me if I am looking at the wrong API, but the `get_xcom` method here:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/execution_api/routes/xcoms.py#L117
Returns only one XCom result by signature and as I understood through code reading:
```python
# The xcom_query allows no map_index to be passed. This endpoint should always return just a single item,
# so we override that query value
xcom_query = xcom_query.filter(BaseXCom.map_index == map_index)
# We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
# This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
# retrieves the raw serialized value from the database. By not relying on `XCom.get_many` or `XCom.get_one`
# (which automatically deserializes using the backend), we avoid potential
# performance hits from retrieving large data files into the API server.
result = xcom_query.limit(1).first()
```
And `XComResult`, `XComResponse` are single-item models, am I missing something? Will it be multiple inside the JSON? My manual tests showed that with the current API call non-providing is not a working strategy, at least back then. I will recheck.
> not passing a map index as "return all XComs for mapped tis"
This behavior is available at `get_xcom_entries` from public xcom API, and I was targeting it in my initial variant, but after testing I understood that different API is used and I cannot rely on this:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/core_api/routes/public/xcom.py#L109
Please, can you guide me, on what kind of behavior is expected from this method? I will try to make it work via `LazySequence` if possible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1986965421
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
> I am sorry, but I don't understand how to utilize the current LazyXComSequence with the arbitrary task_ids.
You don't, it's not designed for that -- only to lazily get the mapped values of a single task.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1986984483
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
So what is the right path to take here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1998056149
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
@amoghrajesh @ashb
Do you think anything else should be done to resolve this thread?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975367422
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -1075,44 +1076,86 @@ def test_get_variable_from_context(
assert var_from_context == Variable(key="test_key", value=expected_value)
+ @pytest.mark.parametrize(
+ "map_indexes",
+ [
+ pytest.param(-1, id="no_map_index"),
+ pytest.param(1, id="single_map_index"),
+ pytest.param(42, id="not_set_with_index"),
+ pytest.param([0, 1], id="multiple_map_indexes"),
+ pytest.param((0, 1), id="any_iterable"),
+ pytest.param(None, id="map_indexes_none"),
+ pytest.param(-9999, id="not_set_no_index"),
+ ],
Review Comment:
I don't think these parameters correctly captures Passing None vs not passing anything at all.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975514391
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -1075,44 +1076,86 @@ def test_get_variable_from_context(
assert var_from_context == Variable(key="test_key", value=expected_value)
+ @pytest.mark.parametrize(
+ "map_indexes",
+ [
+ pytest.param(-1, id="no_map_index"),
+ pytest.param(1, id="single_map_index"),
+ pytest.param(42, id="not_set_with_index"),
+ pytest.param([0, 1], id="multiple_map_indexes"),
+ pytest.param((0, 1), id="any_iterable"),
+ pytest.param(None, id="map_indexes_none"),
+ pytest.param(-9999, id="not_set_no_index"),
+ ],
Review Comment:
Yeah, this code is unclear.
I've updated the params. Hope I've understood your intention right
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2667551079
Hello @insomnes, apologies for the delayed review. Taking a look at it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1962341277
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -270,13 +271,19 @@ def xcom_pull(
:param run_id: If provided, only pulls XComs from a DagRun w/a matching run_id.
If *None* (default), the run_id of the calling task is used.
+
When pulling one single task (``task_id`` is *None* or a str) without
- specifying ``map_indexes``, the return value is inferred from whether
- the specified task is mapped. If not, value from the one single task
- instance is returned. If the task to pull is mapped, an iterator (not a
- list) yielding XComs from mapped task instances is returned. In either
- case, ``default`` (*None* if not specified) is returned if no matching
- XComs are found.
+ specifying ``map_indexes``, the return value is a single XCom entry
+ (map_indexes is set to map_index of the calling task instance).
+
+ When pulling task is mapped the specified ``map_index`` is used, so by default
+ pulling on mapped task will result in no matching XComs if the task instance
+ of the method call is not mapped. Otherwise the map_index of the calling task
+ instance is used. Setting ``map_indexes`` to *None* will pull XCom as from
+ not mapped task.
Review Comment:
Changed as requested
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2734879188
Rebased on recent changes to `task_runner.py`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2728400609
I've resolved two comments, they were properly fixed, I believe. Still left the thread with the API limitations discussion unresolved, because I am unsure if it is appropriate to do it myself in this case. Other threads were simple fixes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2676892986
Up. Would be nice to know what's next.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2657423001
I am still not sure if this is the intended implementation or not, but it is limited by the current execution API restriction of 1 XCom result in response. I didn't touch it to not break anything.
Prints from original DAG example:
```python3
@dag.task
def xcom_pull(**context):
ti: TaskInstance = context["ti"]
# [1, 3, None, None, None, None]"
for x in [
ti.xcom_pull(task_ids=["identity", "foo", "plusfive"], map_indexes=[0, 2]),
]:
print(x)
@dag.task
def xcom_multi_pull(nothing, **context):
ti: TaskInstance = context["ti"]
print("MY TI", ti.map_index)
# For ti 2
# [3, None, None]
for x in [
ti.xcom_pull(task_ids=["identity", "foo", "plusfive"], map_indexes=None),
]:
print(x)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh merged PR #46728:
URL: https://github.com/apache/airflow/pull/46728
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975397441
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
Correct me if I am looking at the wrong API, but the `get_xcom` method here:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/execution_api/routes/xcoms.py#L117
Returns only one XCom result by signature and as I understood through code reading:
```python
# The xcom_query allows no map_index to be passed. This endpoint should always return just a single item,
# so we override that query value
xcom_query = xcom_query.filter(BaseXCom.map_index == map_index)
# We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
# This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
# retrieves the raw serialized value from the database. By not relying on `XCom.get_many` or `XCom.get_one`
# (which automatically deserializes using the backend), we avoid potential
# performance hits from retrieving large data files into the API server.
result = xcom_query.limit(1).first()
```
And `XComResult`, `XComResponse` are single-item models, am I missing something? Will it be multiple inside the JSON? My manual tests showed that with the current API call non-providing is not a working strategy, at least back then. I will recheck.
> not passing a map index as "return all XComs for mapped tis"
This behavior is available at `get_xcom_entries` from public xcom API, and I was targeting it in my initial variant, but after testing I understood that different API is used and I cannot rely on this:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/core_api/routes/public/xcom.py#L109
Can you guide me here, what kind of behavior is expected from this method? I will try to make it work via `LazySequence` if possible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975397441
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
Correct me if I am looking at the wrong API, but the `get_xcom` method here:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/execution_api/routes/xcoms.py#L117
Returns only one XCom result by signature and as I understood through code reading:
```python
# The xcom_query allows no map_index to be passed. This endpoint should always return just a single item,
# so we override that query value
xcom_query = xcom_query.filter(BaseXCom.map_index == map_index)
# We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
# This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
# retrieves the raw serialized value from the database. By not relying on `XCom.get_many` or `XCom.get_one`
# (which automatically deserializes using the backend), we avoid potential
# performance hits from retrieving large data files into the API server.
result = xcom_query.limit(1).first()
```
And `XComResult`, `XComResponse` are single-item models, am I missing something? Will it be multiple inside the JSON? My manual tests showed that with the current API call non-providing is not a working strategy, at least back then. I will recheck.
> not passing a map index as "return all XComs for mapped tis"
This behavior is available at `get_xcom_entries` from public xcom API, and I was targeting it in my initial variant, but after testing I understood that different API is used and I cannot rely on this:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/core_api/routes/public/xcom.py#L109
Can you guide me, please, on what kind of behavior is expected from this method? I will try to make it work via `LazySequence` if possible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2735541107
@insomnes i have rebased it again, once green, will merge, thanks for your patience
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2735656167
> @insomnes i have rebased it again, once green, will merge, thanks for your patience
Thank you, checks are green now.
No problem. I am happy to help with the airflow development!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2663483475
If this PR is not in the right direction I am ok with full rework. I just need to understand what was planned for this feature then.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1998053341
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
Review Comment:
Resolving, feel free to reopen if not
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1992853742
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -275,12 +276,17 @@ def xcom_pull(
If *None* (default), the run_id of the calling task is used.
When pulling one single task (``task_id`` is *None* or a str) without
- specifying ``map_indexes``, the return value is inferred from whether
- the specified task is mapped. If not, value from the one single task
- instance is returned. If the task to pull is mapped, an iterator (not a
- list) yielding XComs from mapped task instances is returned. In either
- case, ``default`` (*None* if not specified) is returned if no matching
- XComs are found.
+ specifying ``map_indexes``, the return value is a single XCom entry
+ (map_indexes is set to map_index of the calling task instance).
+
+ When pulling task is mapped the specified ``map_index`` is used, so by default
+ pulling on mapped task will result in no matching XComs if the task instance
+ of the method call is not mapped. Otherwise the map_index of the calling task
Review Comment:
```suggestion
of the method call is not mapped. Otherwise, the map_index of the calling task
```
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -296,24 +302,48 @@ def xcom_pull(
task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = map_indexes
+ else:
+ raise TypeError(
+ "map_indexes can be omitted or must be an int, an iterable of ints, "
+ f"or None, got {type(map_indexes)}"
+ )
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
Review Comment:
This is meant to go?
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -296,24 +302,48 @@ def xcom_pull(
task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = map_indexes
+ else:
+ raise TypeError(
+ "map_indexes can be omitted or must be an int, an iterable of ints, "
+ f"or None, got {type(map_indexes)}"
+ )
Review Comment:
Yeah with the right indent
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -296,24 +302,48 @@ def xcom_pull(
task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = map_indexes
+ else:
+ raise TypeError(
+ "map_indexes can be omitted or must be an int, an iterable of ints, "
+ f"or None, got {type(map_indexes)}"
+ )
Review Comment:
nit:
```suggestion
raise TypeError(
f"Invalid type for map_indexes: expected int, iterable of ints, or None, got {type(map_indexes)}"
)
```
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -296,24 +302,48 @@ def xcom_pull(
task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = map_indexes
+ else:
+ raise TypeError(
+ "map_indexes can be omitted or must be an int, an iterable of ints, "
+ f"or None, got {type(map_indexes)}"
+ )
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
Review Comment:
```suggestion
# TODO: AIP 72 Execution API only allows working with a single map_index at a time
```
We are adding the AIP 72 as we have done it for multiple places in the repo, easier to grep at a later point
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1960981166
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
Review Comment:
Nice!
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
cc: @ashb as we had some discussion re this in the past
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
Review Comment:
Can we be specific with an elif here?
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -270,13 +271,19 @@ def xcom_pull(
:param run_id: If provided, only pulls XComs from a DagRun w/a matching run_id.
If *None* (default), the run_id of the calling task is used.
+
When pulling one single task (``task_id`` is *None* or a str) without
- specifying ``map_indexes``, the return value is inferred from whether
- the specified task is mapped. If not, value from the one single task
- instance is returned. If the task to pull is mapped, an iterator (not a
- list) yielding XComs from mapped task instances is returned. In either
- case, ``default`` (*None* if not specified) is returned if no matching
- XComs are found.
+ specifying ``map_indexes``, the return value is a single XCom entry
+ (map_indexes is set to map_index of the calling task instance).
+
+ When pulling task is mapped the specified ``map_index`` is used, so by default
+ pulling on mapped task will result in no matching XComs if the task instance
+ of the method call is not mapped. Otherwise the map_index of the calling task
+ instance is used. Setting ``map_indexes`` to *None* will pull XCom as from
+ not mapped task.
Review Comment:
```suggestion
instance is used. Setting ``map_indexes`` to *None* will pull XCom as it would from
a non mapped task.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1962342042
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
Review Comment:
Sure, I also added a specific error on the final else branch, if it's not needed I can drop it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on PR #46728:
URL: https://github.com/apache/airflow/pull/46728#issuecomment-2689925769
Up. If this PR should wait for something else or other solution should be chosen, than it would be nice to have some update about it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975514391
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -1075,44 +1076,86 @@ def test_get_variable_from_context(
assert var_from_context == Variable(key="test_key", value=expected_value)
+ @pytest.mark.parametrize(
+ "map_indexes",
+ [
+ pytest.param(-1, id="no_map_index"),
+ pytest.param(1, id="single_map_index"),
+ pytest.param(42, id="not_set_with_index"),
+ pytest.param([0, 1], id="multiple_map_indexes"),
+ pytest.param((0, 1), id="any_iterable"),
+ pytest.param(None, id="map_indexes_none"),
+ pytest.param(-9999, id="not_set_no_index"),
+ ],
Review Comment:
Yeah, this code is unclear.
I've updated the params. Hope I've understood your idea right
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975397441
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
Correct me if I am looking at the wrong API, but the `get_xcom` method here:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/execution_api/routes/xcoms.py#L117
Returns only one XCom result by signature and as I understood through code reading:
```python
# The xcom_query allows no map_index to be passed. This endpoint should always return just a single item,
# so we override that query value
xcom_query = xcom_query.filter(BaseXCom.map_index == map_index)
# We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
# This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
# retrieves the raw serialized value from the database. By not relying on `XCom.get_many` or `XCom.get_one`
# (which automatically deserializes using the backend), we avoid potential
# performance hits from retrieving large data files into the API server.
result = xcom_query.limit(1).first()
```
And `XComResult`, `XComResponse` are single-item models, am I missing something? Will it be multiple inside the JSON? My manual tests showed that with the current API call non-providing is not a working strategy, at least back then. I will recheck.
> not passing a map index as "return all XComs for mapped tis"
This behavior is available at `get_xcom_entries` from public xcom API, and I was targeting it in my initial variant, but after testing I understood that different API is used and I cannot rely on this:
https://github.com/apache/airflow/blob/main/airflow/api_fastapi/core_api/routes/public/xcom.py#L109
Please, can you guide me on what kind of behavior is expected from this method? I will try to make it work via `LazySequence` if possible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1975364946
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
So yeah, I think the endpoint will handle not passing a map index as "return all XComs for mapped tis". So this comment isn't right, and LazySequence now exists.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ability to pull XCom when passing multiple map_indexes [airflow]
Posted by "insomnes (via GitHub)" <gi...@apache.org>.
insomnes commented on code in PR #46728:
URL: https://github.com/apache/airflow/pull/46728#discussion_r1986132894
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -289,27 +296,46 @@ def xcom_pull(
if task_ids is None:
# default to the current task if not provided
- task_ids = self.task_id
+ task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
+
+ map_indexes_iterable: Iterable[int | None] = []
+ # If map_indexes is not provided, default to use the map_index of the calling task
if isinstance(map_indexes, ArgNotSet):
- map_indexes = self.map_index
- elif isinstance(map_indexes, Iterable):
- # TODO: Handle multiple map_indexes or remove support
- raise NotImplementedError("Multiple map_indexes are not supported yet")
+ map_indexes_iterable = [self.map_index]
+ elif isinstance(map_indexes, int) or map_indexes is None:
+ map_indexes_iterable = [map_indexes]
+ else:
+ map_indexes_iterable = map_indexes
log = structlog.get_logger(logger_name="task")
xcoms = []
- for t in task_ids:
+ # TODO: Execution API only allows working with a single map_index at a time
+ # this is inefficient and leads to task_id * map_index requests to the API.
+ # And we can't achieve the original behavior of XCom pull with multiple tasks
+ # directly now.
+ # Original behavior may be achieved after `LazyXComSequence` is finished?
+ #
+ # Original description:
+ #
+ # When pulling one single task (``task_id`` is *None* or a str) without
+ # specifying ``map_indexes``, the return value is inferred from whether
+ # the specified task is mapped. If not, value from the one single task
+ # instance is returned. If the task to pull is mapped, an iterator (not a
+ # list) yielding XComs from mapped task instances is returned. In either
+ # case, ``default`` (*None* if not specified) is returned if no matching
+ # XComs are found.
+ for t_id, m_idx in product(task_ids, map_indexes_iterable):
Review Comment:
@ashb
I am sorry, but I don't understand how to utilize the current `LazyXComSequence` with the arbitrary task_ids. Because it requires `RuntimeTaskInstance` and `PlainXComArg` it's pulling and I don't get a convenient way to get another task `RuntimeTaskInstance` instance and `PlainXComArg` / `Operator` in this method.
Maybe it should be changed to use `SUPERVISOR_COMMS` for `xcom_pull` and operate on `GetXCom` / `GetXComCount` fields directly and not on `PlainXComArg` / `RuntimeTaskInstance`.
Or I could create another similar class for this case.
Otherwise, some other execution API endpoint should be introduced not limited by `xcom_query.limit(1).first()`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org