pierrejeambrun commented on code in PR #76297: URL: https://github.com/apache/airflow/pull/76297#discussion_r4481603684
########## airflow-core/adr/lang-sdk/0016-mapped-tasks-across-lang-sdks.md: ########## @@ -0,0 +1,440 @@ +<!-- + 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. + --> + +# ADR-0016: Mapped Tasks Across Lang SDKs + +## Status + +Proposed. + +## Terms + +| Term | Means | +|---|---| +| Expansion | Turning one mapped task into N task instances while the Dag run is in flight. | +| Expand input | The serialized description of what a task expands over. One per mapped task. | +| Map length | How many instances an expansion produces. | +| Map index | An instance's position in its expansion, `0..N-1`. `-1` is the placeholder that has not expanded yet. | +| Region | The row that scopes one expansion, so the same mapped task expands independently per enclosing iteration. | +| Producer | The task whose output is expanded over. | +| Consumer | The mapped task itself. | +| Binding | One resolved call argument delivered to a runtime. | +| Reduce | Reading a mapped task's whole output from a downstream task, as a sequence rather than one value. | + +## Decision + +1. **Expansion stays in the scheduler for every language.** Map length, instance creation, regions, `REMOVED` handling and trigger rules are core's, not an SDK's. + An SDK emits an expand input and consumes one value per map index; it never decides how many instances exist. +2. **The expand input keeps Python's `{__type, __var}` encoding.** An SDK emits the same form the Dag processor already writes for a Python mapped task, so core parses one format and no SDK invents a second. +3. **Map length is computed once in the supervisor, instead of being re-implemented in each SDK.** The server tells the producer's supervisor whether the length is needed; the supervisor measures the return value and rejects one that is absent or unsized. + The API server loads no extra Dag, and Python's own copy is deleted so the two cannot drift. +4. **Every Python mapping primitive is in scope, and the server resolves all of them.** A primitive an SDK has not implemented yet is a capability flag, not a core limit. +5. **A runtime always consumes delivered bindings; it never resolves the expand input itself.** The server picks this index's element in both authoring modes, so no SDK re-implements element selection or cross-product arithmetic. +6. **Only `map()` callables run outside the server, applied on top of the element it picked.** Selection and transformation separate: [`_MapResult`](https://github.com/apache/airflow/blob/27b8c82f54d44830176eb049f2e42916c7849bb0/task-sdk/src/airflow/sdk/definitions/xcom_arg.py#L398) indexes first, then applies the callables. + A Dag definition's callables belong to its own SDK and are applied by that runtime, which reads them from its own bundle. A Python Dag's callables are Python, so the helper applies them and the binding arrives as a plain literal. + No language reads another's callable. +7. **The map-index label is a runtime call, not a template.** `map_index_template` is Jinja and no other language renders it; an SDK sends `SetRenderedMapIndex` instead. +8. **Reduce is a lazy sequence, never a materialized value.** A downstream consumer reads length and elements on demand, as [`LazyXComSequence`](https://github.com/apache/airflow/blob/27b8c82f54d44830176eb049f2e42916c7849bb0/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py#L64) does. + +## Context + +Mapped task behaviour spans four layers: the authoring API builds it, serialization carries it, +the scheduler materializes it, and the worker resolves one element. Python implements all four. An +SDK that copied that split would re-implement the scheduler's half in its own language, and the two +would diverge. + +An SDK owns *what to expand over* and *what to do with one element*, and core owns *how many* and +*which*. The same split holds inside a task loop: a region scopes one expansion to one iteration, +so a mapped task in a loop expands once per iteration with no extra work in the SDK. Task loops +therefore need no mapping design of their own. + +## Who runs what + +The coordinator is the supervisor of a Lang-SDK task, so the two names mean one process. + +| Process | Language | Its share of mapping | +|---|---|---| +| Scheduler | Python | Counts the expansion and turns the `-1` placeholder into N instances. | +| API server, Execution API | Python | Picks this index's element and builds the bindings for `ti_run`. | +| Supervisor / coordinator | Python, in the worker process | Measures the producer's map length; forwards every XCom request. | +| `map()` helper | Python, one per consuming task instance | Imports a Python Dag file and applies its `map()` callables. | +| SDK runtime | the SDK's own language | Runs the handler, applies the SDK's own `map()` callables, holds the reduce sequence. | + +## Authoring interface + +- `expand` takes the arguments that vary per instance; `partial` takes the arguments that stay the same at every instance. +- A mapped task returns the same reference type as any other task, so an edge or a later `expand` treats it identically. +- `zip`, `concat` and `map` are methods on a reference, not free functions, so they chain in declaration order. +- A reduce consumer receives a sequence of the element type, not the element type. + +The running example, in the TypeScript SDK. Every later section follows this same `extract → +transform → report`. + +```ts +async function extract(): Promise<string[]>; // ["uk ", " fr", "de "] +async function transform(args: { region: string; row: string }): Promise<number>; +async function report(args: { totals: XComSequence<number> }): Promise<void>; +``` + +```ts +const rows = dag.task("extract", extract); +const totals = dag.task("transform", transform).partial({ region: "eu" }).expand({ row: rows }); +dag.task("report", report, { totals }); +``` + +`transform` runs once per element of `extract`'s output and receives one `row`; `report` runs once +and receives every result, as decision 8's lazy sequence. + +## What the Dag serialization carries + +The expand input holds the items themselves for a literal, and only a reference for another task's +output. A literal therefore needs no producer: the scheduler creates every instance with the Dag +run. + +```jsonc +// .expand({ row: ["uk ", " fr", "de "] }) +"expand_input": {"type": "dict-of-lists", "value": {"__type": "dict", "__var": { + "row": ["uk ", " fr", "de "]}}} + +// .expand({ row: rows }) +"expand_input": {"type": "dict-of-lists", "value": {"__type": "dict", "__var": { + "row": {"__type": "xcomref", "__var": {"task_id": "extract", "key": "return_value"}}}}} + +// .expand({ row: rows.map(clean) }) +"expand_input": {"type": "dict-of-lists", "value": {"__type": "dict", "__var": { + "row": {"__type": "xcomref", "__var": { + "type": "map", + "arg": {"task_id": "extract", "key": "return_value"}, + "callables": ["(row) => row.trim()"]}}}}} + +// .expand({ row: rows.zip(lookup, { fillvalue: "" }) }) +"expand_input": {"type": "dict-of-lists", "value": {"__type": "dict", "__var": { + "row": {"__type": "xcomref", "__var": { + "type": "zip", + "args": [{"task_id": "extract", "key": "return_value"}, + {"task_id": "lookup", "key": "return_value"}], + "fillvalue": ""}}}}} +``` + +Two details an SDK emitter has to get right: + +- The outer `type` is the expand input's, `dict-of-lists` or `list-of-dicts`. The inner `type` is the XCom argument's, one of `map`, `zip` and `concat`; a plain argument omits it. +- Only the outermost argument is wrapped in `xcomref`. A nested one — `map`'s `arg`, `zip`'s and `concat`'s `args` — is the bare dict, because [`serialize_xcom_arg`](https://github.com/apache/airflow/blob/27b8c82f54d44830176eb049f2e42916c7849bb0/task-sdk/src/airflow/sdk/definitions/xcom_arg.py#L590) recurses without re-wrapping. + +A `map()` callable is stored as source text, one string per callable in call order: Python writes +`inspect.getsource(f)`, an SDK writes its own equivalent — `fn.toString()` in TypeScript. Core never +deserializes these strings; they exist so the UI can show the transform. Execution always re-enters Review Comment: Small one: I don't think anything in the UI reads `callables` today. Maybe "for display only", as the Consequences section puts 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
