dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r3501034354
##########
task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py:
##########
@@ -79,6 +79,103 @@ def _needs_run_time_resolution(v: OperatorExpandArgument)
-> TypeGuard[MappedArg
return isinstance(v, (MappedArgument, XComArg))
+def count(expand_input: ExpandInput, iterable: Iterable[Any]) -> Iterable[Any]:
+ expand_input._length = None
+ counter = 0
+
+ for item in iterable:
+ counter += 1
+ yield item
+
+ expand_input._length = counter
+
+
[email protected]()
+class ExpandInput(ABC, ResolveMixin):
+ EXPAND_INPUT_TYPE: ClassVar[str]
+ _length: int | None
+
+ def __attrs_post_init__(self):
+ self._length = None
+
+ @property
+ @abstractmethod
+ def value(self) -> Any:
+ """The value of the expand input."""
+ ...
+
+ def iter_values(self, context: Mapping[str, Any]) -> Iterable[Any]:
+ raise NotImplementedError()
+
+ def resolve(self, context: Mapping[str, Any]) -> Any:
+ raise NotImplementedError()
+
+ def __len__(self) -> int:
+ if self._length is None:
+ raise RuntimeError(f"Length of {type(self).__name__} is not yet
known")
+ return self._length
+
+
+class DecoratedExpandInput(ExpandInput):
+ EXPAND_INPUT_TYPE: ClassVar[str] = "decorated"
+
+ def __init__(self, expand_input: ExpandInput):
+ super().__init__()
+ self.delegate = expand_input
+
+ @property
+ def value(self) -> Any:
+ return self.delegate.value
+
+ def iter_references(self) -> Iterable[tuple[Operator, str]]:
+ return self.delegate.iter_references()
+
+ def iter_values(self, context: Mapping[str, Any]) -> Iterable[dict]:
+ return map(
+ lambda value: {"op_kwargs": value},
+ self.delegate.iter_values(context),
+ )
+
+ def resolve(self, context: Mapping[str, Any]) -> tuple[Mapping[str, Any],
set[int]]:
+ return self.delegate.resolve(context)
+
+ def __len__(self) -> int:
+ return len(self.delegate)
+
+
+class BatchedExpandInput(DecoratedExpandInput):
+ """
+ ExpandInput that batches another ExpandInput into N chunks.
+
+ This affects mapping cardinality, NOT resolve-time behavior.
+ """
+
+ EXPAND_INPUT_TYPE: ClassVar[str] = "batched"
+
+ def __init__(self, expand_input: ExpandInput, size: int):
+ super().__init__(expand_input=expand_input)
+ self.size = size
+
+ def iter_values(self, context: Mapping[str, Any]) -> Iterable[dict]:
+ map_index = context["ti"].map_index
+
+ for index, item in enumerate(self.delegate.iter_values(context)):
+ if index % self.size == map_index:
+ yield item
+
+ def __len__(self) -> int:
+ total = len(self.delegate)
+
+ base = total // self.size
+ remainder = total % self.size
+
+ # each map_index corresponds to a specific bucket size
+ # but we don't know which instance this is without context
+
+ # so we return worst-case or max-case:
+ return base + (1 if remainder > 0 else 0)
Review Comment:
I re-did the same logic with the count on iter_values.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
