dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r4028338120
########## task-sdk/src/airflow/sdk/definitions/batchedoperator.py: ########## @@ -0,0 +1,525 @@ +# +# 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. +from __future__ import annotations + +import inspect +from abc import ABCMeta, abstractmethod +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +import attrs + +from airflow.sdk import TriggerRule, timezone +from airflow.sdk.bases.decorator import ( + DecoratedMappedOperator, + FParams, + FReturn, + OperatorSubclass, + _TaskDecorator, + get_unique_task_id, +) +from airflow.sdk.bases.operator import ( + BaseOperator, + coerce_resources, + coerce_timedelta, + get_merged_defaults, + parse_retries, +) +from airflow.sdk.definitions._internal.contextmanager import ( + DagContext, + TaskGroupContext, +) +from airflow.sdk.definitions._internal.expandinput import ( + EXPAND_INPUT_EMPTY, + DecoratedExpandInput, + DictOfListsExpandInput, + ExpandInput, + ListOfDictsExpandInput, + OperatorExpandArgument, + OperatorExpandKwargsArgument, +) +from airflow.sdk.definitions._internal.types import NOTSET +from airflow.sdk.definitions.mappedoperator import ( + MappedOperator, + OperatorPartial, + ensure_xcomarg_return_value, + prevent_duplicates, + validate_mapping_kwargs, +) +from airflow.sdk.definitions.xcom_arg import XComArg + +if TYPE_CHECKING: + from airflow.sdk.definitions.iterableoperator import IterableOperator, MappedIterableOperator + from airflow.sdk.definitions.mappedoperator import ValidationSource + from airflow.sdk.definitions.param import ParamsDict + +T = TypeVar("T", bound=OperatorPartial | _TaskDecorator) + + [email protected](kw_only=True, repr=False) +class BatchableOperator(Generic[T], metaclass=ABCMeta): + """ + Intermediate abstraction for batched mapping. + + This class decorates an OperatorPartial and stores configuration for batched mapping. + It is used to facilitate batched expansion of operators, allowing tasks to be mapped over batches + of data and then iterate over the batched data. + + :param operator_partial: The partial operator to be batched. + :param size: The number of task instances to create. The input is distributed across them + round-robin (item ``i`` goes to task instance ``i % size``), not split into ``size`` + contiguous chunks — this is *not* the same semantics as ``itertools.batched(iterable, size)``. + See :class:`~airflow.sdk.definitions._internal.expandinput.BatchedExpandInput` for why + round-robin is used instead of contiguous chunking. + """ + + operator_partial: T + size: int + + @property + def operator_class(self) -> type[BaseOperator]: + return self.operator_partial.operator_class + + @property + def kwargs(self) -> dict[str, Any]: + return self.operator_partial.kwargs + + @abstractmethod + def iterate(self, **mapped_kwargs: OperatorExpandArgument) -> Any: + """ + Iterate the operator over the provided mapped keyword arguments. + + :param mapped_kwargs: Keyword arguments to expand against. + :return: An expanded operator or XComArg, depending on the subclass implementation. + """ + + @abstractmethod + def iterate_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True) -> Any: + """ + Iterate the operator over a list of dictionaries or XComArg. + + :param kwargs: List of dicts or XComArg to expand against. + :param strict: Whether to enforce strict argument checking. + :return: An expanded operator or XComArg, depending on the subclass implementation. + """ + + @abstractmethod + def _iterate( + self, + expand_input: ExpandInput, + *, + strict: bool, + ) -> IterableOperator | MappedIterableOperator: + """ + Create an iterable operator for the given expansion input. + + This method calls the _expand method first to get a MappedOperator based on expansion input, + then wraps it in either an IterableOperator or MappedIterableOperator depending on the batch size. + + :param expand_input: The input to iterate against. + :param strict: Whether to enforce strict argument checking. + :return: An IterableOperator or MappedIterableOperator. + """ + + @abstractmethod + def _expand( + self, + expand_input: ExpandInput, + *, + strict: bool, + register_with_dag: bool = True, + ) -> MappedOperator: + """ + Create a mapped operator for the given expansion input. + + :param expand_input: The input to expand against. + :param strict: Whether to enforce strict argument checking. + :param register_with_dag: Whether to apply upstream relationships. + :return: A MappedOperator instance. + """ + + [email protected](kw_only=True, repr=False) +class BatchedOperator(BatchableOperator[OperatorPartial]): + """ + Concrete implementation of BatchableOperator for classic (non-decorated) operators. + + This class wraps an OperatorPartial and provides batched expansion and iteration logic + for classic Airflow operators. It enables mapping tasks over batches of data, supporting + both direct expansion via keyword arguments and expansion via a list of dictionaries or XComArg. + + :param operator_partial: The OperatorPartial instance to be batched and expanded. + :param size: The number of task instances to create for mapping. Items are distributed across + them round-robin (item ``i`` goes to task instance ``i % size``), not split into ``size`` + contiguous chunks. + """ + + @property + def params(self) -> ParamsDict | dict: + return self.operator_partial.params + + @property + def _expand_called(self) -> bool: + return self.operator_partial._expand_called + + @_expand_called.setter + def _expand_called(self, value: bool) -> None: + self.operator_partial._expand_called = value + + def iterate(self, **mapped_kwargs: OperatorExpandArgument) -> IterableOperator | MappedIterableOperator: + if not mapped_kwargs: + raise TypeError("no arguments to iterate against") + + validate_mapping_kwargs(self.operator_class, "iterate", mapped_kwargs) + prevent_duplicates( + self.kwargs, + mapped_kwargs, + fail_reason="unmappable or already specified", + ) + # Since the input is already checked at parse time, we can set strict + # to False to skip the checks on execution. + expand_input = DictOfListsExpandInput(mapped_kwargs) + return self._iterate(expand_input, strict=False) + + def iterate_kwargs( + self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True + ) -> IterableOperator | MappedIterableOperator: + if isinstance(kwargs, Sequence): + for item in kwargs: + if not isinstance(item, (XComArg, Mapping)): + raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") + elif not isinstance(kwargs, XComArg): + raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") + + expand_input = ListOfDictsExpandInput(kwargs) + return self._iterate(expand_input, strict=strict) + + def _iterate( + self, + expand_input: ExpandInput, + *, + strict: bool, + ) -> IterableOperator | MappedIterableOperator: + from airflow.sdk.definitions.iterableoperator import IterableOperator, MappedIterableOperator + + # Unlike .expand(), neither OperatorPartial.iterate()/.iterate_kwargs() nor this class's own + # iterate()/iterate_kwargs() set _expand_called, so OperatorPartial.__del__ would otherwise + # warn "Task ... was never mapped!" even though .iterate()/.batch().iterate() legitimately + # consumed the partial. + self._expand_called = True + operator = self._expand(expand_input, strict=strict, register_with_dag=False) + + if self.size > 1: Review Comment: Fixed in 784bc54195. `batch()` now rejects any size below two with a `ValueError`, on both the classic and the decorated path. The reason it was unvalidated is that the public method also served as the internal funnel: `.iterate()` and `.expand()` call it with a `size=0` sentinel so the batched and unbatched paths share one code path. That funnel is now a private `_batch()`, so the sentinel is unreachable from user code and the public method can validate. `BatchedExpandInput`'s own check stays as the runtime guard and now agrees with it. --- Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting -- 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]
