jason810496 opened a new pull request, #70571: URL: https://github.com/apache/airflow/pull/70571
- **Depends on https://github.com/apache/airflow/pull/70209 and https://github.com/apache/airflow/pull/70570 get merged first** - **Diff for early review**: https://github.com/jason810496/airflow/compare/feature/go-sdk/taskflow-arg-binding...feature/go-sdk/taskflow-arg-binding-mapped ## Why The mapped lang-SDK PR makes the execution API derive a **per-map-index** arg-binding spec for a mapped `@task.stub`. This PR makes the Go SDK actually consume it: a stub expanded with `.expand()` delivers *its* element to the native Go function, and a Go stub's output can feed a downstream `.expand()`. ## Supported dynamic-mapping forms A mapped `@task.stub` produces **one arg-binding per map index** (declaration order), and the Go runtime binds each onto the native function's parameters. Each form and how the runtime delivers it: **Expanded arguments** — `.expand(param=…)`, value differs per index: | Form | Example | How the Go runtime binds it | |---|---|---| | Expand over a **literal collection** | `via.expand(country=["uk", "fr", "de"])` | server resolves the element; arrives as a **plain literal** | | Expand over an **unmapped upstream's list output** | `via.expand(extracted=extract())` | pull the unmapped XCom row, then **take `element_index`** of the list | | Expand over a **mapped upstream** | `via.expand(extracted=seed.expand(n=[1, 2]))` | pull the upstream XCom **row at `map_index`** directly | **Multiple expanded arguments** — cross product: | Form | Example | How the Go runtime binds it | |---|---|---| | `.expand(a=…, b=…)` | `combine.expand(a=["x", "y"], b=[1, 2, 3])` → 6 instances | the server decomposes the map index per kwarg; the runtime binds whatever each kwarg resolved to, so expand sources can be **mixed in one call** | **Partial arguments** — `.partial(param=…)`, constant across every index: | Form | Example | How the Go runtime binds it | |---|---|---| | Partial **literal** | `via.partial(region="uk").expand(…)` | plain literal, same value every index | | Partial **unmapped-upstream XCom** | `via.partial(config=load_config()).expand(…)` | pull the whole return value, same every index | **Defaulted arguments:** | Form | Example | How the Go runtime binds it | |---|---|---| | Unpassed param with a **signature default** | `retries: int = 3` left unpassed | literal `from_default`; keyword-style `sdk.TaskInput` fields may leave it unclaimed | One DAG exercising every form, and the Go function receiving it: ```python @task.stub(queue="golang") def make_items(): ... # unmapped stub → returns a list @task.stub(queue="golang") def seed(n: int): ... # mapped below → a mapped upstream @task.stub(queue="golang") def transform( country: str, # expand over a literal list extracted: dict, # expand over an unmapped upstream's list → element_index seeded: dict, # expand over a mapped upstream → map_index region: str, # partial literal (constant) config: dict, # partial xcom over an unmapped upstream (constant) retries: int = 3, # unpassed → default captured (from_default) ): ... @dag(dag_id="mapped_binding_dag") def mapped_binding_dag(): transform.partial( region="uk", config=load_config(), ).expand( # cross product; expand sources mixed across kwargs country=["uk", "fr", "de"], extracted=make_items(), seeded=seed.expand(n=[1, 2, 3]), ) ``` ```go // One instance per (country, extracted-element, seeded-row) combination. Flat // params bind in declaration order after the injectables; every expanded arg // arrives already narrowed to this instance's element. func Transform( ctx sdk.TIRunContext, log *slog.Logger, country string, extracted map[string]any, seeded map[string]any, region string, config map[string]any, retries int, ) (any, error) ``` **Rejected loudly** (server-side, mirroring the provider's parse-time checks): `.expand_kwargs()`; a `partial()` kwarg over a **mapped** upstream's aggregated output; `.map()` / `.zip()` / `concat` or custom-key XCom; non-JSON literals; a mapped stub TI still at `map_index=-1`. A mapped stub in an **older-provider** Dag delivers no bindings and keeps the legacy ignored-args behavior. ## How - **Delivery (`pkg/binding`):** `XComArg` gains `MapIndex` / `ElementIndex` (from regenerated `genmodels`). `Resolve` now pulls the specific upstream row `MapIndex` selects (expand over a mapped upstream); when `ElementIndex` is set it takes that element of the pulled sequence (expand over an unmapped upstream's list output), with typed/out-of-range errors. An unmapped argument carries neither and takes the whole value; literal expands are resolved to their element server-side and arrive as plain literals — so the flat-vs-struct binding surface from #70209 is unchanged. - **Return side (`mapped_length`):** a foreign runtime can't inspect the Dag to learn its return value feeds a downstream `.expand()`. New server-derived `TIRunContext.has_mapped_dependants` flag (computed from `iter_mapped_dependants`); when set, the supervisor records `mapped_length = len(value)` on the return-value `SetXCom` on the task's behalf — the foreign-runtime analogue of the Python task runner's `_push_xcom_if_needed` logic — so the scheduler can expand the Go stub's mapped dependants. - **Server:** the single serialized-Dag lookup now returns `(arg_bindings, has_mapped_dependants)`; both fields are version-gated to `2026-10-30` so pre-arg-bindings clients skip the derivation entirely. --- ##### Was generative AI tooling used to co-author this PR? - [x] Yes, with help of Claude Code (Opus 4.8) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) -- 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]
