apache / apache/airflow

Deliver value_schema on arg bindings for mapped stub tasks

Open
#70,523 0 comments 0 reactions 1 assignee Claimed by @jason810496 View on GitHub
area:coordinator area:core area:go-sdk area:java-sdk area:ts-sdk kind:feature on hold
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 9h
Merged PRs (30d)
472

Description

### Background

https://github.com/apache/airflow/pull/69757 lets `@task.stub` tasks be called TaskFlow-style: a parse-time capture builds the arg-binding spec (parameter `name`, `kind`, `value`, and a pydantic-generated `value_schema` JSON-schema fragment inferred from the stub function's annotations) and ships it to lang-SDK runtimes via `TIRunContext.arg_bindings`.

Mapped stubs (`.expand()`) never instantiate `_StubOperator` at parse time — `_TaskDecorator._expand()` builds a generic `DecoratedMappedOperator` — so no spec lands in the serialized Dag. Instead ti_run derives per-map-index bindings server-side in `airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py` (`_resolve_mapped_stub_arg_bindings`). The stub function's annotations are not available on the API server, so derived bindings omit `value_schema` and foreign runtimes fall back to decode-only checks.

### What needs to happen

1. Add a capture point at `.expand()` time in the task-sdk decorator machinery (operator-class hook in `_TaskDecorator._expand()`, or at mapped-operator serialization time in the Dag processor, where the callable is importable) so the stub's per-parameter schemas (`_infer_value_schema` in `providers/standard/src/airflow/providers/standard/decorators/stub.py`) ride with the mapped operator.
2. Serialize the captured schema map on mapped operators (new serialized field + `airflow-core/src/airflow/serialization/schema.json` entry).
3. Merge the schemas into the server-derived bindings in `_resolve_mapped_stub_arg_bindings` so mapped bindings carry `value_schema` like unmapped ones.
4. Degrade gracefully on Airflow 2.x installs of the standard provider (no hook/pydantic → bindings omit `value_schema`, today's behavior).
5. Tests: task-sdk expand capture, Dag serialization round trip, execution API ti_run response.

### Acceptance criteria

- A mapped `@task.stub` with annotated parameters receives `value_schema` on each derived binding in `TIRunContext.arg_bindings`.
- The limitation note + tracking link in `_resolve_mapped_stub_arg_bindings` are removed.
- The omission contract (absent key = unconstrained) and unmapped behavior are unchanged.

### Context

- Originating PR: https://github.com/apache/airflow/pull/69757
- Review thread: https://github.com/apache/airflow/pull/69757#discussion_r3654919186
- Parent effort: #66937

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.