dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r3992218352
########## task-sdk/src/airflow/sdk/definitions/batchedoperator.py: ########## @@ -0,0 +1,512 @@ +# +# 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 batches to create. Review Comment: Good catch — size was ambiguous. Clarified in the docstrings (BatchableOperator, BatchedOperator, DecoratedBatchedOperator, and BatchedExpandInput): size is the number of task instances to create, and items are distributed round-robin (index % size == map_index), not into size contiguous chunks. On the choice itself: round-robin vs. contiguous chunking aren't actually different in how lazily they can filter items — index % size and index // size are both pure functions of index and size, neither needs N to route a given item. The real constraint is that Airflow has to fix the number of task instances before any of them run (get_parse_time_mapped_ti_count/get_mapped_ti_count). With round-robin, that count is size itself — a constant the user picks, independent of N. The only place N enters is the min(batch_size, mapped_ti_count) clamp to avoid empty TIs when N < size, and that only needs a bounded lookahead (or already-tracked mapped-length metadata), not the true final count. With a chunk-length semantic (size = items per instance, à la itertools.batched), the TI count would be ceil(N / size) — which is undefined until N is fully known. For the streamed/paginated-API sources this feature targets, that means draining the whole input just to size the mapping, which defeats the point. So round-robin isn't an arbitrary pick — it's the only one of the two where size can be fixed without first knowing N. Happy to switch to chunking for a future .chunk(size=N) variant that explicitly documents its "materializes/counts the input first" trade-off, but for .batch()/.iterate() I'd keep round-robin. -- 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]
