You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "uranusjr (via GitHub)" <gi...@apache.org> on 2024/06/11 09:37:28 UTC
[PR] Implement ChainXComArg [airflow]
uranusjr opened a new pull request, #40172:
URL: https://github.com/apache/airflow/pull/40172
This is useful when you want to do a thing against more than one list of things. You can sort of already do this with an extra task, but this is lazier and saves some resources both in the XCom storage and memory needed to run tasks.
--
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] Implement XComArg concat() [airflow]
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #40172:
URL: https://github.com/apache/airflow/pull/40172#discussion_r1642835149
##########
airflow/models/xcom_arg.py:
##########
@@ -606,8 +615,78 @@ def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
return _ZipResult(values, fillvalue=self.fillvalue)
+class _ConcatResult(Sequence):
+ def __init__(self, values: Sequence[Sequence | dict]) -> None:
+ self.values = values
+
+ def __getitem__(self, index: Any) -> Any:
+ i = index
+ for value in self.values:
+ if i >= (curlen := len(value)):
+ i -= curlen
+ elif isinstance(value, Sequence):
+ return value[i]
+ else:
+ return next(itertools.islice(iter(value), i, None))
+ raise IndexError(index)
+
+ def __len__(self) -> int:
+ return sum(len(v) for v in self.values)
+
+
+class ConcatXComArg(XComArg):
+ """Concatenating multiple XCom references into one.
+
+ This is done by calling ``concat()`` on an XComArg to combine it with
+ others. The effect is similar to Python's :func:`itertools.chain`, but the
+ return value also supports index access.
+ """
+
+ def __init__(self, args: Sequence[XComArg]) -> None:
+ if not args:
+ raise ValueError("At least one input is required")
+ self.args = args
+
+ def __repr__(self) -> str:
+ args_iter = iter(self.args)
+ first = repr(next(args_iter))
+ rest = ", ".join(repr(arg) for arg in args_iter)
+ return f"{first}.concat({rest})"
+
+ def _serialize(self) -> dict[str, Any]:
+ return {"args": [serialize_xcom_arg(arg) for arg in self.args]}
+
+ @classmethod
+ def _deserialize(cls, data: dict[str, Any], dag: DAG) -> XComArg:
+ return cls([deserialize_xcom_arg(arg, dag) for arg in data["args"]])
+
+ def iter_references(self) -> Iterator[tuple[Operator, str]]:
+ for arg in self.args:
+ yield from arg.iter_references()
+
+ def concat(self, *others: XComArg) -> ConcatXComArg:
+ # Flattern foo.concat(x).concat(y) into one call.
Review Comment:
```suggestion
# Flatten foo.concat(x).concat(y) into one call.
```
Small typo here.
Relatedly, would it also be useful to explicitly call out this functionality in the documentation too?
--
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] Implement ChainXComArg [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2164435040
> I briefly searched how other ecosystems call this and `concat` seems popular (behind `chain`). `combine` is more often like `join` but I think it still is an option in this case since the joining semantic does not make sense.
Maybe `append` is a better name ? I agree with @hussein-awala -> we need some docs to make it discoverable, it's a rather useful feature, but having a description and example of use is pretty much a "MUST HAVE" 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] Implement ChainXComArg [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2172234222
I think I’m going to go with `concat`. This is used in [JavaScript](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Array/concat). Java has `Stream.concat` (which is more like `itertools.chain`, close enough). Various languages also have `concat` for strings. Probably among the most solid choices beside `chain`.
--
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] Implement ChainXComArg [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2172729248
SGTM
--
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] Implement XComArg concat() [airflow]
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2173804067
> > concat XComArgs from two sets of mapped tasks
>
> Not sure I follow this
Related to a question I could think of DAG authors asking: "Does this work with mapped-task inputs?".
```mermaid
graph LR;
mapped_task_1-->aggregate_mapped_tasks;
mapped_task_2-->aggregate_mapped_tasks;
```
meaning
```python
aggregate_mapped_tasks(input=mapped_task_1.concat(mapped_task_2))
```
Looking again the "... ``concat`` function takes arbitrary positional arguments ..." statement covers this so ignore. Still looks great.
--
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] Implement ChainXComArg [airflow]
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2164294012
Oddly enough, this seems useful for the discussion in #40124?
--
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] Implement XComArg concat() [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2174920741
I added some text to call out chaining calls.
--
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] Implement ChainXComArg [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2165328765
I don’t like `append` personally since in Python it’s a method that modifies the object in-place, not creating a new object. (Same for `extend` although nobody mentions it yet.) I plan to add docs after we agree on the interface.
--
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] Implement XComArg concat() [airflow]
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #40172:
URL: https://github.com/apache/airflow/pull/40172#discussion_r1642985445
##########
airflow/models/xcom_arg.py:
##########
@@ -606,8 +615,78 @@ def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
return _ZipResult(values, fillvalue=self.fillvalue)
+class _ConcatResult(Sequence):
+ def __init__(self, values: Sequence[Sequence | dict]) -> None:
+ self.values = values
+
+ def __getitem__(self, index: Any) -> Any:
+ i = index
+ for value in self.values:
+ if i >= (curlen := len(value)):
+ i -= curlen
+ elif isinstance(value, Sequence):
+ return value[i]
+ else:
+ return next(itertools.islice(iter(value), i, None))
+ raise IndexError(index)
+
+ def __len__(self) -> int:
+ return sum(len(v) for v in self.values)
+
+
+class ConcatXComArg(XComArg):
+ """Concatenating multiple XCom references into one.
+
+ This is done by calling ``concat()`` on an XComArg to combine it with
+ others. The effect is similar to Python's :func:`itertools.chain`, but the
+ return value also supports index access.
+ """
+
+ def __init__(self, args: Sequence[XComArg]) -> None:
+ if not args:
+ raise ValueError("At least one input is required")
+ self.args = args
+
+ def __repr__(self) -> str:
+ args_iter = iter(self.args)
+ first = repr(next(args_iter))
+ rest = ", ".join(repr(arg) for arg in args_iter)
+ return f"{first}.concat({rest})"
+
+ def _serialize(self) -> dict[str, Any]:
+ return {"args": [serialize_xcom_arg(arg) for arg in self.args]}
+
+ @classmethod
+ def _deserialize(cls, data: dict[str, Any], dag: DAG) -> XComArg:
+ return cls([deserialize_xcom_arg(arg, dag) for arg in data["args"]])
+
+ def iter_references(self) -> Iterator[tuple[Operator, str]]:
+ for arg in self.args:
+ yield from arg.iter_references()
+
+ def concat(self, *others: XComArg) -> ConcatXComArg:
+ # Flattern foo.concat(x).concat(y) into one call.
Review Comment:
I was just thinking of questions authors would have. Like "can I 'chain' the `.concat()` calls?" or "even can I pass a list of XComArgs in one call e.g.,`foo.concat([x, y])`?". The docs say you can use multiple iterables, but a trivial mention might go a long way.
--
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] Implement XComArg concat() [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #40172:
URL: https://github.com/apache/airflow/pull/40172#discussion_r1642862399
##########
airflow/models/xcom_arg.py:
##########
@@ -606,8 +615,78 @@ def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
return _ZipResult(values, fillvalue=self.fillvalue)
+class _ConcatResult(Sequence):
+ def __init__(self, values: Sequence[Sequence | dict]) -> None:
+ self.values = values
+
+ def __getitem__(self, index: Any) -> Any:
+ i = index
+ for value in self.values:
+ if i >= (curlen := len(value)):
+ i -= curlen
+ elif isinstance(value, Sequence):
+ return value[i]
+ else:
+ return next(itertools.islice(iter(value), i, None))
+ raise IndexError(index)
+
+ def __len__(self) -> int:
+ return sum(len(v) for v in self.values)
+
+
+class ConcatXComArg(XComArg):
+ """Concatenating multiple XCom references into one.
+
+ This is done by calling ``concat()`` on an XComArg to combine it with
+ others. The effect is similar to Python's :func:`itertools.chain`, but the
+ return value also supports index access.
+ """
+
+ def __init__(self, args: Sequence[XComArg]) -> None:
+ if not args:
+ raise ValueError("At least one input is required")
+ self.args = args
+
+ def __repr__(self) -> str:
+ args_iter = iter(self.args)
+ first = repr(next(args_iter))
+ rest = ", ".join(repr(arg) for arg in args_iter)
+ return f"{first}.concat({rest})"
+
+ def _serialize(self) -> dict[str, Any]:
+ return {"args": [serialize_xcom_arg(arg) for arg in self.args]}
+
+ @classmethod
+ def _deserialize(cls, data: dict[str, Any], dag: DAG) -> XComArg:
+ return cls([deserialize_xcom_arg(arg, dag) for arg in data["args"]])
+
+ def iter_references(self) -> Iterator[tuple[Operator, str]]:
+ for arg in self.args:
+ yield from arg.iter_references()
+
+ def concat(self, *others: XComArg) -> ConcatXComArg:
+ # Flattern foo.concat(x).concat(y) into one call.
Review Comment:
The functionality is exactly the same from the user’s perspective if this isn’t done. Calling out such optimisation seems awkward to me.
--
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] Implement XComArg concat() [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr merged PR #40172:
URL: https://github.com/apache/airflow/pull/40172
--
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] Implement XComArg concat() [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2172961431
Alright, I’ve made the name change and added docs.
--
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] Implement ChainXComArg [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2166670508
Maybe just `xcom_chain()` or `chain_xcom()` ? That would remove confusion and keep itertools relation - at the expense of somewhat redundancy.
--
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] Implement ChainXComArg [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2164332423
I briefly searched how other ecosystems call this and `concat` seems popular (behind `chain`). `combine` is more often like `join` but I think it still is an option in this case since the joining semantic does not make sense.
--
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] Implement XComArg concat() [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2172977385
> This is useful when you want to do a thing against more than one list of things.
Very nice! BTW I think it's also interesting to be able to do thing against specific items in the XCOM rather than all of 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] Implement XComArg concat() [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #40172:
URL: https://github.com/apache/airflow/pull/40172#issuecomment-2173470913
> concat XComArgs from two sets of mapped tasks
Not sure I follow this
--
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