jason810496 opened a new pull request, #70570: URL: https://github.com/apache/airflow/pull/70570
- **Depends on https://github.com/apache/airflow/pull/69757 get merged first** - **Diff for early review**: https://github.com/jason810496/airflow/compare/feature/lang-sdk/taskflow-stub-dag...feature/lang-sdk/taskflow-stub-dag-mapped related: lifts the `.expand()`/`.partial()` restriction #69757 declared out of scope; the Go SDK runtime that consumes this is the stacked follow-up. ## Why #69757 ships stub TaskFlow arg-binding but **rejects `.expand()` on a stub at parse time**. This PR restores that scope: a `@task.stub` can be dynamically task-mapped, and every map index gets its own arg-binding spec so the foreign runtime receives *its* element instead of the aggregated output. ## Supported dynamic-mapping forms A mapped `@task.stub` produces **one arg-binding per map index**, always in the stub's **signature declaration order** (not call-site order). Each parameter resolves as one of: **Expanded arguments** — `.expand(param=…)`, value differs per index: | Form | Example | Wire binding | |---|---|---| | Expand over a **literal collection** | `transform.expand(country=["uk", "fr", "de"])` | `literal`, element resolved server-side (list → element; dict → `[key, value]` per item) | | Expand over an **unmapped upstream's list output** | `transform.expand(extracted=extract())` | `xcom` + `element_index=i` (pull the single row, take element `i`) | | Expand over a **mapped upstream** | `transform.expand(extracted=seed.expand(n=[1, 2]))` | `xcom` + `map_index=i` (pull upstream row `i` directly) | **Multiple expanded arguments** — cross product: | Form | Example | Behavior | |---|---|---| | `.expand(a=…, b=…)` | `combine.expand(a=["x", "y"], b=[1, 2, 3])` → 6 instances | the map index is decomposed into one sub-index per kwarg (last varies fastest); each kwarg resolves independently, so the three expand forms above can be **mixed in one call** | **Partial arguments** — `.partial(param=…)`, constant across every index: | Form | Example | Wire binding | |---|---|---| | Partial **literal** | `transform.partial(country="uk").expand(…)` | `literal` (same value every index) | | Partial **unmapped-upstream XCom** | `transform.partial(extracted=extract()).expand(…)` | `xcom`, whole return value (no sub-index) | **Defaulted arguments:** | Form | Example | Wire binding | |---|---|---| | Unpassed param with a **signature default** | `retries: int = 3` left unpassed | `literal` + `from_default: true` | One DAG exercising every form at once: ```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]), ) ``` **Rejected loudly** (parse-time in the provider; re-checked server-side for Dags from other provider versions): `.expand_kwargs()`; a `partial()` kwarg over a **mapped** upstream's aggregated output (would bind the nonexistent `map_index=-1` row); `.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 (no captured metadata) delivers no bindings and keeps the legacy ignored-args behavior. ## How - **Parse time (providers/standard):** a mapped stub never instantiates at parse time, so `_StubOperator` captures per-parameter metadata (declaration order, defaults, value schemas) via a new optional `get_mapped_serialized_fields` operator hook. The core serializer calls it at the single point where `operator_class`/`python_callable` are still the real objects — everything the server cannot recover from the serialized Dag alone. - **Wire model:** `XComArgBinding` regains `map_index` (which upstream row to pull — expand over a *mapped* upstream) and `element_index` (take element N of the unmapped list — expand over an *unmapped* upstream's output). - **Server (`ti_run`):** for a mapped stub, bindings are derived per map index by decomposing the TI's `map_index` into one sub-index per expanded kwarg via new `SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes` — the server-side twin of the SDK's `_expand_mapped_field` cross-product (last kwarg varies fastest). Expanded kwargs get `map_index`/`element_index`, `partial()` kwargs and unpassed defaults bind as above. - **Backward compat:** Dags serialized by an older provider carry no metadata → resolve to `None` → keep the legacy ignored-args behavior (their args were never deliverable). The provider's parse-time rejections are re-checked server-side for Dags produced by other provider versions. - **Schema:** new `arg_binding_param` definition + optional `_mapped_arg_binding_params` array on the operator; the inner object stays open so newer metadata keeps validating on older cores. No `SERIALIZER_VERSION` bump (optional field). --- ##### 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]
